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:
    20→
    21→/// Heartbeat timeout (90 seconds - 3x the agent's 30 second interval)
    22→const HEARTBEAT_TIMEOUT_SECS: u64 = 90;
    23→
    24→/// Session state
    25→#[derive(Debug, Clone)]
    26→pub struct Session {
    27→    pub id: SessionId,
    28→    pub agent_id: AgentId,
    29→    pub agent_name: String,
    30→    pub started_at: chrono::DateTime<chrono::Utc>,
    31→    pub viewer_count: usize,
    32→    pub is_streaming: bool,
    33→    pub is_online: bool,  // Whether agent is currently connected
    34→    pub is_persistent: bool,  // Persistent agent (no support code) vs support session
    35→    pub last_heartbeat: chrono::DateTime<chrono::Utc>,
    36→    // Agent status info
    37→    pub os_version: Option<String>,
    38→    pub is_elevated: bool,
    39→    pub uptime_secs: i64,
    40→    pub display_count: i32,
    41→}
    42→
    43→/// Channel for sending frames from agent to viewers
    44→pub type FrameSender = broadcast::Sender<Vec<u8>>;
    45→pub type FrameReceiver = broadcast::Receiver<Vec<u8>>;

/// Session state
#[derive(Debug, Clone)]
pub struct Session {
    pub id: SessionId,
    pub agent_id: AgentId,
    pub agent_name: String,
    pub started_at: chrono::DateTime<chrono::Utc>,
    pub viewer_count: usize,
    pub is_streaming: bool,
    pub 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,
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        (session_id, frame_tx, input_rx)
    }

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

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

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

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

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

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

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

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

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

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

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

        Some((frame_rx, input_tx))
    }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

+    pub is_online: bool,  // Whether agent is currently connected

+    pub is_persistent: bool,  // Persistent agent (no support code) vs support session


Now I need to update the register_agent function to:
1. Check if an existing offline session exists for this agent_id
2. If so, reuse it and mark it online
3. Add is_persistent parameter
4. Also need to update remove_session to mark offline instead of removing for persistent agentsLet me continue editing.


Edit

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

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

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

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

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

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

        (session_id, frame_tx, input_rx)
    }

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


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:
    76→            agents: Arc::new(RwLock::new(HashMap::new())),
    77→        }
    78→    }
    79→
    80→    /// Register a new agent and create a session
    81→    /// If agent was previously connected (offline session exists), reuse that session
    82→    pub async fn register_agent(&self, agent_id: AgentId, agent_name: String, is_persistent: bool) -> (SessionId, FrameSender, InputReceiver) {
    83→        // Check if this agent already has an offline session (reconnecting)
    84→        {
    85→            let agents = self.agents.read().await;
    86→            if let Some(&existing_session_id) = agents.get(&agent_id) {
    87→                let mut sessions = self.sessions.write().await;
    88→                if let Some(session_data) = sessions.get_mut(&existing_session_id) {
    89→                    if !session_data.info.is_online {
    90→                        // Reuse existing session - mark as online and create new channels
    91→                        tracing::info!("Agent {} reconnecting to existing session {}", agent_id, existing_session_id);
    92→
    93→                        let (frame_tx, _) = broadcast::channel(16);
    94→                        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64);
    95→
    96→                        session_data.info.is_online = true;
    97→                        session_data.info.last_heartbeat = chrono::Utc::now();
    98→                        session_data.info.agent_name = agent_name; // Update name in case it changed
    99→                        session_data.frame_tx = frame_tx.clone();
   100→                        session_data.input_tx = input_tx;
   101→                        session_data.last_heartbeat_instant = Instant::now();
   102→
   103→                        return (existing_session_id, frame_tx, input_rx);
   104→                    }
   105→                }
   106→            }
   107→        }
   108→
   109→        // Create new session
   110→        let session_id = Uuid::new_v4();
   111→
   112→        // Create channels
   113→        let (frame_tx, _) = broadcast::channel(16); // Buffer 16 frames
   114→        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64); // Buffer 64 input events
   115→
   116→        let now = chrono::Utc::now();
   117→        let session = Session {
   118→            id: session_id,
   119→            agent_id: agent_id.clone(),
   120→            agent_name,
   121→            started_at: now,
   122→            viewer_count: 0,
   123→            is_streaming: false,
   124→            is_online: true,
   125→            is_persistent,
   126→            last_heartbeat: now,
   127→            os_version: None,
   128→            is_elevated: false,
   129→            uptime_secs: 0,
   130→            display_count: 1,
   131→        };
   132→
   133→        let session_data = SessionData {
   134→            info: session,
   135→            frame_tx: frame_tx.clone(),
   136→            input_tx,
   137→            input_rx: None,
   138→            viewers: HashSet::new(),
   139→            last_heartbeat_instant: Instant::now(),
   140→        };
   141→
   142→        let mut sessions = self.sessions.write().await;
   143→        sessions.insert(session_id, session_data);
   144→
   145→        let mut agents = self.agents.write().await;
   146→        agents.insert(agent_id, session_id);
   147→
   148→        (session_id, frame_tx, input_rx)
   149→    }
   150→
   151→    /// Update agent status from heartbeat or status message
   152→    pub async fn update_agent_status(
   153→        &self,

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

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

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

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

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

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

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

/// Session state
#[derive(Debug, Clone)]
pub struct Session {
    pub id: SessionId,
    pub agent_id: AgentId,
    pub agent_name: String,
    pub started_at: chrono::DateTime<chrono::Utc>,
    pub viewer_count: usize,
    pub is_streaming: bool,
    pub 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
    pub async fn register_agent(&self, agent_id: AgentId, agent_name: String) -> (SessionId, FrameSender, InputReceiver) {
        let session_id = Uuid::new_v4();

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

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

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

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

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

        (session_id, frame_tx, input_rx)
    }

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

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

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

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

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

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

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

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

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

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

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

        Some((frame_rx, input_tx))
    }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

+    /// 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);

started_at: now,
viewer_count: 0,
is_streaming: false,
+            is_online: true,
+            is_persistent,
last_heartbeat: now,
os_version: None,
is_elevated: false,


Now I need to update remove_session to mark offline instead of removing for persistent agents. Let me find and update that function.


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


Conversation compacted


This session is being continued from a previous conversation that ran out of context. The conversation is summarized below:
Analysis:
Let me chronologically analyze this conversation:

1. **Session Context Recovery**: The conversation started with a summary from a previous session about GuruConnect - a remote desktop solution. Key infrastructure was already in place with a server at connect.azcomputerguru.com.

2. **User's First Request**: "My concern is both zombie sessions as well as what time of performance/bandwidth will be used if I have 2000 machines connected?"
   - I explained the issue: agents were streaming immediately upon connection at 30 FPS
   - Proposed idle/active mode architecture

3. **User's Second Request**: "yes, implement idle/active mode"
   - I implemented StartStream/StopStream/AgentStatus protocol messages
   - Rewrote agent session module with idle/streaming states
   - Updated server to track viewers and send start/stop commands
   - Added heartbeat mechanism with 90-second timeout
   - Fixed a compile error (Self:: vs self. for static methods)
   - Deployed to server

4. **User's Third Request**: "I installed the agent on another laptop. What can we test?"
   - Showed the agent ACG-TECH-01L connected in idle mode
   - Reviewed server logs showing successful start/stop stream transitions

5. **User's Fourth Request**: "Let's build a viewer"
   - Entered plan mode
   - Explored codebase with subagents
   - Created plan for vanilla JS viewer
   - Implemented viewer.html with:
     - Canvas-based video display
     - Protobuf parsing for VideoFrame/RawFrame
     - Zstd decompression via fzstd CDN
     - BGRA to RGBA conversion
     - Mouse/keyboard event capture and encoding
     - Toolbar with disconnect, fullscreen, Ctrl+Alt+Del
   - Updated dashboard.html to open viewer
   - Deployed to server

6. **User's Fifth Request**: "/s" - Save session
   - Updated session log with all work done
   - Committed and pushed

7. **User's Sixth (Current) Request**: "That laptop was asleep, and disappeard from sessions. It should have remained in sessionlist but show as offline when sleeping?"
   - Started implementing is_online/is_persistent flags
   - Updated Session struct with new fields
   - Updated register_agent to handle reconnection
   - Was in the middle of updating remove_session to mark offline instead of removing

Key files modified:
- proto/guruconnect.proto - Added StartStream, StopStream, AgentStatus
- agent/src/session/mod.rs - Idle/active state machine
- agent/src/main.rs - Pass is_elevated to SessionManager
- server/src/session/mod.rs - Viewer tracking, heartbeat, is_online/is_persistent
- server/src/relay/mod.rs - Handle new messages
- server/src/api/mod.rs - Added new fields to API
- server/static/viewer.html - New file, full viewer implementation
- server/static/dashboard.html - Open viewer instead of alert

Errors encountered:
1. "no method named `send_stop_stream_internal` found" - Fixed by using Self:: instead of self.

Current work: Adding is_online flag to sessions so persistent agents show as "offline" when disconnected instead of disappearing. I've updated Session struct and register_agent, need to update remove_session.

Summary:
1. Primary Request and Intent:
   The user is building GuruConnect, a ScreenConnect replacement for remote desktop support. In this session, the explicit requests were:
   - Implement idle/active mode for scalable agent connections (concern about 2000 agents)
   - Build a remote desktop viewer
   - Fix sessions to show "offline" instead of disappearing when agent disconnects (current task)

2. Key Technical Concepts:
   - Idle/Active mode: Agents connect in idle mode (heartbeat only), transition to streaming when viewer connects
   - StartStream/StopStream protocol messages for controlling agent state
   - Heartbeat mechanism with 90-second timeout for zombie detection
   - Protobuf encoding/decoding in JavaScript (manual varint parsing)
   - Zstd decompression in browser using fzstd library
   - BGRA to RGBA conversion for HTML5 canvas
   - Persistent vs support session distinction (is_persistent flag)
   - Session reconnection handling for offline agents

3. Files and Code Sections:

   - **proto/guruconnect.proto** - Added new messages for idle/active mode
     ```protobuf
     message StartStream {
         string viewer_id = 1;
         int32 display_id = 2;
     }
     message StopStream {
         string viewer_id = 1;
     }
     message AgentStatus {
         string hostname = 1;
         string os_version = 2;
         bool is_elevated = 3;
         int64 uptime_secs = 4;
         int32 display_count = 5;
         bool is_streaming = 6;
     }
     ```

   - **agent/src/session/mod.rs** - Complete rewrite with idle/streaming state machine
     - Lazy initialization of capturer/encoder only when streaming
     - Heartbeat every 30s, status every 60s in idle mode
     - Handles StartStream/StopStream messages

   - **server/src/session/mod.rs** - CURRENTLY BEING MODIFIED
     - Added is_online and is_persistent fields to Session
     - Updated register_agent to handle reconnection:
     ```rust
     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
                         session_data.info.is_online = true;
                         // ... create new channels ...
                         return (existing_session_id, frame_tx, input_rx);
                     }
                 }
             }
         }
         // Create new session with is_online: true, is_persistent
     }
     ```
     - Need to update remove_session to mark offline instead of removing

   - **server/static/viewer.html** - NEW FILE, 684 lines
     - Canvas-based video display with protobuf parsing
     - Mouse/keyboard input encoding
     - Zstd decompression, BGRA→RGBA conversion

   - **server/static/dashboard.html** - Updated joinSession() and connectToMachine() to open viewer

4. Errors and fixes:
   - **"no method named `send_stop_stream_internal` found for reference `&SessionManager`"**:
     - Cause: Used `self.send_stop_stream_internal()` for a static method
     - Fix: Changed to `Self::send_stop_stream_internal()`
     - Committed fix: cc35d11

5. Problem Solving:
   - Solved scalability concern: Agents now use ~100 bytes/min in idle mode vs 100+ KB/sec when streaming
   - Solved viewer implementation: Full canvas viewer with input support
   - Currently solving: Agent sessions disappearing when laptop sleeps (should show offline)

6. All user messages:
   - "My concern is both zombie sessions as well as what time of performance/bandwidth will be used if I have 2000 machines connected?"
   - "yes, implement idle/active mode"
   - "I installed the agent on another laptop. What can we test?"
   - "Let's build a viewer"
   - "/s" (save session)
   - "That laptop was asleep, and disappeard from sessions. It should have remained in sessionlist but show as offline when sleeping?"

7. Pending Tasks:
   - Complete is_online/is_persistent implementation:
     - Update remove_session to mark offline instead of removing (for persistent agents)
     - Update relay/mod.rs to pass is_persistent to register_agent
     - Update API response to include is_online
     - Update dashboard to show online/offline status
   - Test viewer with live agent (viewer was just deployed)

8. Current Work:
   Implementing online/offline status tracking for persistent agents. The Session struct and register_agent have been updated with is_online and is_persistent fields. The next step is to update remove_session at line 291.

   User said: "That laptop was asleep, and disappeard from sessions. It should have remained in sessionlist but show as offline when sleeping?"

   I had just run grep to find remove_session function:
   ```
   291:    pub async fn remove_session(&self, session_id: SessionId) {
   ```

9. Optional Next Step:
   Update `remove_session` in server/src/session/mod.rs to mark sessions as offline (is_online = false) instead of removing them for persistent agents. Then update relay/mod.rs to pass `is_persistent = support_code.is_none()` to register_agent, update the API, and update the dashboard UI.

   Direct quote from user's most recent request: "That laptop was asleep, and disappeard from sessions. It should have remained in sessionlist but show as offline when sleeping?".
Please continue the conversation from where we left it off without asking the user any further questions. Continue with the last task that you were asked to work on.


The user wants me to continue implementing the is_online/is_persistent feature. Looking at the context:

1. I've already updated Session struct with is_online and is_persistent fields
2. I've updated register_agent to handle reconnection
3. I need to:
   - Update remove_session to mark offline instead of removing for persistent agents
   - Update relay/mod.rs to pass is_persistent to register_agent
   - Update API response to include is_online
   - Update dashboard to show online/offline status

Let me continue where I left off. I need to:
1. First, update remove_session in session/mod.rs
2. Then update relay/mod.rs
3. Then update api/mod.rs
4. Then update dashboard.html

Let me read the current state of these files to continue the work.


Continuing with the is_online/is_persistent implementation. I need to update `remove_session` to mark sessions offline instead of removing them for persistent agents.


   280→                viewer_id: viewer_id.to_string(),
   281→            })),
   282→        };
   283→
   284→        let mut buf = Vec::new();
   285→        if stop_stream.encode(&mut buf).is_ok() {
   286→            let _ = session_data.input_tx.send(buf).await;
   287→        }
   288→    }
   289→
   290→    /// Remove a session (when agent disconnects)
   291→    pub async fn remove_session(&self, session_id: SessionId) {
   292→        let mut sessions = self.sessions.write().await;
   293→        if let Some(session_data) = sessions.remove(&session_id) {
   294→            let mut agents = self.agents.write().await;
   295→            agents.remove(&session_data.info.agent_id);
   296→        }
   297→    }
   298→
   299→    /// Disconnect a session by sending a disconnect message to the agent
   300→    /// Returns true if the message was sent successfully
   301→    pub async fn disconnect_session(&self, session_id: SessionId, reason: &str) -> bool {
   302→        let sessions = self.sessions.read().await;
   303→        if let Some(session_data) = sessions.get(&session_id) {
   304→            // Create disconnect message
   305→            use crate::proto;
   306→            use prost::Message;
   307→
   308→            let disconnect_msg = proto::Message {
   309→                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
   310→                    reason: reason.to_string(),
   311→                })),
   312→            };
   313→
   314→            let mut buf = Vec::new();
   315→            if disconnect_msg.encode(&mut buf).is_ok() {
   316→                // Send via input channel (will be forwarded to agent's WebSocket)
   317→                if session_data.input_tx.send(buf).await.is_ok() {
   318→                    return true;
   319→                }
   320→            }
   321→        }
   322→        false
   323→    }
   324→
   325→    /// List all active sessions
   326→    pub async fn list_sessions(&self) -> Vec<Session> {
   327→        let sessions = self.sessions.read().await;
   328→        sessions.values().map(|s| s.info.clone()).collect()
   329→    }
   330→}
   331→
   332→impl Default for SessionManager {
   333→    fn default() -> Self {
   334→        Self::new()
   335→    }
   336→}
   337→

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

                viewer_id: viewer_id.to_string(),
            })),
        };

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

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

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

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

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

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

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


     1→//! WebSocket relay handlers
     2→//!
     3→//! Handles WebSocket connections from agents and viewers,
     4→//! relaying video frames and input events between them.
     5→
     6→use axum::{
     7→    extract::{
     8→        ws::{Message, WebSocket, WebSocketUpgrade},
     9→        Query, State,
    10→    },
    11→    response::IntoResponse,
    12→};
    13→use futures_util::{SinkExt, StreamExt};
    14→use prost::Message as ProstMessage;
    15→use serde::Deserialize;
    16→use tracing::{error, info, warn};
    17→use uuid::Uuid;
    18→
    19→use crate::proto;
    20→use crate::session::SessionManager;
    21→use crate::AppState;
    22→
    23→#[derive(Debug, Deserialize)]
    24→pub struct AgentParams {
    25→    agent_id: String,
    26→    #[serde(default)]
    27→    agent_name: Option<String>,
    28→    #[serde(default)]
    29→    support_code: Option<String>,
    30→    #[serde(default)]
    31→    hostname: Option<String>,
    32→}
    33→
    34→#[derive(Debug, Deserialize)]
    35→pub struct ViewerParams {
    36→    session_id: String,
    37→}
    38→
    39→/// WebSocket handler for agent connections
    40→pub async fn agent_ws_handler(
    41→    ws: WebSocketUpgrade,
    42→    State(state): State<AppState>,
    43→    Query(params): Query<AgentParams>,
    44→) -> impl IntoResponse {
    45→    let agent_id = params.agent_id;
    46→    let agent_name = params.hostname.or(params.agent_name).unwrap_or_else(|| agent_id.clone());
    47→    let support_code = params.support_code;
    48→    let sessions = state.sessions.clone();
    49→    let support_codes = state.support_codes.clone();
    50→
    51→    ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, agent_id, agent_name, support_code))
    52→}
    53→
    54→/// WebSocket handler for viewer connections
    55→pub async fn viewer_ws_handler(
    56→    ws: WebSocketUpgrade,
    57→    State(state): State<AppState>,
    58→    Query(params): Query<ViewerParams>,
    59→) -> impl IntoResponse {
    60→    let session_id = params.session_id;
    61→    let sessions = state.sessions.clone();
    62→
    63→    ws.on_upgrade(move |socket| handle_viewer_connection(socket, sessions, session_id))
    64→}
    65→
    66→/// Handle an agent WebSocket connection
    67→async fn handle_agent_connection(
    68→    socket: WebSocket,
    69→    sessions: SessionManager,
    70→    support_codes: crate::support_codes::SupportCodeManager,
    71→    agent_id: String,
    72→    agent_name: String,
    73→    support_code: Option<String>,
    74→) {
    75→    info!("Agent connected: {} ({})", agent_name, agent_id);
    76→
    77→    let (mut ws_sender, mut ws_receiver) = socket.split();
    78→
    79→    // If a support code was provided, check if it's valid
    80→    if let Some(ref code) = support_code {
    81→        // Check if the code is cancelled or invalid
    82→        if support_codes.is_cancelled(code).await {
    83→            warn!("Agent tried to connect with cancelled code: {}", code);
    84→            // Send disconnect message to agent
    85→            let disconnect_msg = proto::Message {
    86→                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
    87→                    reason: "Support session was cancelled by technician".to_string(),
    88→                })),
    89→            };
    90→            let mut buf = Vec::new();
    91→            if prost::Message::encode(&disconnect_msg, &mut buf).is_ok() {
    92→                let _ = ws_sender.send(Message::Binary(buf.into())).await;
    93→            }
    94→            let _ = ws_sender.close().await;
    95→            return;
    96→        }
    97→    }
    98→
    99→    // Register the agent and get channels
   100→    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone()).await;
   101→
   102→    info!("Session created: {} (agent in idle mode)", session_id);
   103→
   104→    // If a support code was provided, mark it as connected
   105→    if let Some(ref code) = support_code {
   106→        info!("Linking support code {} to session {}", code, session_id);
   107→        support_codes.mark_connected(code, Some(agent_name.clone()), Some(agent_id.clone())).await;
   108→        support_codes.link_session(code, session_id).await;
   109→    }
   110→
   111→    // Use Arc<Mutex> for sender so we can use it from multiple places
   112→    let ws_sender = std::sync::Arc::new(tokio::sync::Mutex::new(ws_sender));
   113→    let ws_sender_input = ws_sender.clone();
   114→    let ws_sender_cancel = ws_sender.clone();
   115→
   116→    // Task to forward input events from viewers to agent
   117→    let input_forward = tokio::spawn(async move {
   118→        while let Some(input_data) = input_rx.recv().await {
   119→            let mut sender = ws_sender_input.lock().await;
   120→            if sender.send(Message::Binary(input_data.into())).await.is_err() {
   121→                break;
   122→            }
   123→        }
   124→    });
   125→
   126→    let sessions_cleanup = sessions.clone();
   127→    let sessions_status = sessions.clone();
   128→    let support_codes_cleanup = support_codes.clone();
   129→    let support_code_cleanup = support_code.clone();
   130→    let support_code_check = support_code.clone();
   131→    let support_codes_check = support_codes.clone();
   132→
   133→    // Task to check for cancellation every 2 seconds
   134→    let cancel_check = tokio::spawn(async move {
   135→        let mut interval = tokio::time::interval(std::time::Duration::from_secs(2));
   136→        loop {
   137→            interval.tick().await;
   138→            if let Some(ref code) = support_code_check {
   139→                if support_codes_check.is_cancelled(code).await {
   140→                    info!("Support code {} was cancelled, disconnecting agent", code);
   141→                    // Send disconnect message
   142→                    let disconnect_msg = proto::Message {
   143→                        payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
   144→                            reason: "Support session was cancelled by technician".to_string(),
   145→                        })),
   146→                    };
   147→                    let mut buf = Vec::new();
   148→                    if prost::Message::encode(&disconnect_msg, &mut buf).is_ok() {
   149→                        let mut sender = ws_sender_cancel.lock().await;
   150→                        let _ = sender.send(Message::Binary(buf.into())).await;
   151→                        let _ = sender.close().await;
   152→                    }
   153→                    break;
   154→                }
   155→            }
   156→        }
   157→    });
   158→
   159→    // Main loop: receive messages from agent
   160→    while let Some(msg) = ws_receiver.next().await {
   161→        match msg {
   162→            Ok(Message::Binary(data)) => {
   163→                // Try to decode as protobuf message
   164→                match proto::Message::decode(data.as_ref()) {
   165→                    Ok(proto_msg) => {
   166→                        match &proto_msg.payload {
   167→                            Some(proto::message::Payload::VideoFrame(_)) => {
   168→                                // Broadcast frame to all viewers (only sent when streaming)
   169→                                let _ = frame_tx.send(data.to_vec());
   170→                            }
   171→                            Some(proto::message::Payload::ChatMessage(chat)) => {
   172→                                // Broadcast chat message to all viewers
   173→                                info!("Chat from client: {}", chat.content);
   174→                                let _ = frame_tx.send(data.to_vec());
   175→                            }
   176→                            Some(proto::message::Payload::AgentStatus(status)) => {
   177→                                // Update session with agent status
   178→                                sessions_status.update_agent_status(
   179→                                    session_id,
   180→                                    Some(status.os_version.clone()),
   181→                                    status.is_elevated,
   182→                                    status.uptime_secs,
   183→                                    status.display_count,
   184→                                    status.is_streaming,
   185→                                ).await;
   186→                                info!("Agent status update: {} - streaming={}, uptime={}s",
   187→                                    status.hostname, status.is_streaming, status.uptime_secs);
   188→                            }
   189→                            Some(proto::message::Payload::Heartbeat(_)) => {
   190→                                // Update heartbeat timestamp
   191→                                sessions_status.update_heartbeat(session_id).await;
   192→                            }
   193→                            Some(proto::message::Payload::HeartbeatAck(_)) => {
   194→                                // Agent acknowledged our heartbeat
   195→                                sessions_status.update_heartbeat(session_id).await;
   196→                            }
   197→                            _ => {}
   198→                        }
   199→                    }
   200→                    Err(e) => {
   201→                        warn!("Failed to decode agent message: {}", e);
   202→                    }
   203→                }
   204→            }
   205→            Ok(Message::Close(_)) => {
   206→                info!("Agent disconnected: {}", agent_id);
   207→                break;
   208→            }
   209→            Ok(Message::Ping(data)) => {
   210→                // Pong is handled automatically by axum
   211→                let _ = data;
   212→            }
   213→            Ok(_) => {}
   214→            Err(e) => {
   215→                error!("WebSocket error from agent {}: {}", agent_id, e);
   216→                break;
   217→            }
   218→        }
   219→    }
   220→
   221→    // Cleanup
   222→    input_forward.abort();
   223→    cancel_check.abort();
   224→    sessions_cleanup.remove_session(session_id).await;
   225→
   226→    // Mark support code as completed if one was used (unless cancelled)
   227→    if let Some(ref code) = support_code_cleanup {
   228→        if !support_codes_cleanup.is_cancelled(code).await {
   229→            support_codes_cleanup.mark_completed(code).await;
   230→            info!("Support code {} marked as completed", code);
   231→        }
   232→    }
   233→
   234→    info!("Session {} ended", session_id);
   235→}
   236→
   237→/// Handle a viewer WebSocket connection
   238→async fn handle_viewer_connection(
   239→    socket: WebSocket,
   240→    sessions: SessionManager,
   241→    session_id_str: String,
   242→) {
   243→    // Parse session ID
   244→    let session_id = match uuid::Uuid::parse_str(&session_id_str) {
   245→        Ok(id) => id,
   246→        Err(_) => {
   247→            warn!("Invalid session ID: {}", session_id_str);
   248→            return;
   249→        }
   250→    };
   251→
   252→    // Generate unique viewer ID
   253→    let viewer_id = Uuid::new_v4().to_string();
   254→
   255→    // Join the session (this sends StartStream to agent if first viewer)
   256→    let (mut frame_rx, input_tx) = match sessions.join_session(session_id, viewer_id.clone()).await {
   257→        Some(channels) => channels,
   258→        None => {
   259→            warn!("Session not found: {}", session_id);
   260→            return;
   261→        }
   262→    };
   263→
   264→    info!("Viewer {} joined session: {}", viewer_id, session_id);
   265→
   266→    let (mut ws_sender, mut ws_receiver) = socket.split();
   267→
   268→    // Task to forward frames from agent to this viewer
   269→    let frame_forward = tokio::spawn(async move {
   270→        while let Ok(frame_data) = frame_rx.recv().await {
   271→            if ws_sender.send(Message::Binary(frame_data.into())).await.is_err() {
   272→                break;
   273→            }
   274→        }
   275→    });
   276→
   277→    let sessions_cleanup = sessions.clone();
   278→    let viewer_id_cleanup = viewer_id.clone();
   279→
   280→    // Main loop: receive input from viewer and forward to agent
   281→    while let Some(msg) = ws_receiver.next().await {
   282→        match msg {
   283→            Ok(Message::Binary(data)) => {
   284→                // Try to decode as protobuf message
   285→                match proto::Message::decode(data.as_ref()) {
   286→                    Ok(proto_msg) => {
   287→                        match &proto_msg.payload {
   288→                            Some(proto::message::Payload::MouseEvent(_)) |
   289→                            Some(proto::message::Payload::KeyEvent(_)) |
   290→                            Some(proto::message::Payload::SpecialKey(_)) => {
   291→                                // Forward input to agent
   292→                                let _ = input_tx.send(data.to_vec()).await;
   293→                            }
   294→                            Some(proto::message::Payload::ChatMessage(chat)) => {
   295→                                // Forward chat message to agent
   296→                                info!("Chat from technician: {}", chat.content);
   297→                                let _ = input_tx.send(data.to_vec()).await;
   298→                            }
   299→                            _ => {}
   300→                        }
   301→                    }
   302→                    Err(e) => {
   303→                        warn!("Failed to decode viewer message: {}", e);
   304→                    }
   305→                }
   306→            }
   307→            Ok(Message::Close(_)) => {
   308→                info!("Viewer {} disconnected from session: {}", viewer_id, session_id);
   309→                break;
   310→            }
   311→            Ok(_) => {}
   312→            Err(e) => {
   313→                error!("WebSocket error from viewer {}: {}", viewer_id, e);
   314→                break;
   315→            }
   316→        }
   317→    }
   318→
   319→    // Cleanup (this sends StopStream to agent if last viewer)
   320→    frame_forward.abort();
   321→    sessions_cleanup.leave_session(session_id, &viewer_id_cleanup).await;
   322→    info!("Viewer {} left session: {}", viewer_id_cleanup, session_id);
   323→}
   324→

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


Now I have a clear picture. I need to:

1. Update `remove_session` in session/mod.rs to mark persistent sessions as offline instead of removing them
2. Update `register_agent` call in relay/mod.rs to pass `is_persistent` (based on support_code being None)
3. Update API response to include is_online
4. Update dashboard to show online/offline status

Let me start:

1. First, I see that `register_agent` is called at line 100 in relay/mod.rs, and it currently only takes 2 arguments:
```rust
let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone()).await;
```

But I updated register_agent to take 3 arguments (including is_persistent). I need to update the call.

2. The `remove_session` function needs to be changed to mark persistent sessions as offline instead of removing them.

Let me update these files.


TodoWrite

Add is_online flag to sessions

Update remove_session to mark offline (not remove) for persistent agents

Update relay to pass is_persistent to register_agent

Update API response to include is_online

Update dashboard to show online/offline status


Todos have been modified successfully. Ensure that you continue to use the todo list to track your progress. Please proceed with the current tasks if applicable

Add is_online flag to sessions

Keep persistent agent sessions on disconnect (mark offline)

Update dashboard to show online/offline status

Handle agent reconnection (reuse existing session)

Add is_online flag to sessions

Update remove_session to mark offline (not remove) for persistent agents

Update relay to pass is_persistent to register_agent

Update API response to include is_online


Now let me update the remove_session function to mark persistent sessions as offline instead of removing them.


Edit

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

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