From 1cc94c61e791378550dd7164a3a38fa31ac89e4d Mon Sep 17 00:00:00 2001
From: Mike Swanson <mike@azcomputerguru.com>
Date: Sun, 28 Dec 2025 17:52:26 -0700
Subject: [PATCH] Add is_online/is_persistent for persistent agent sessions

- Sessions now track whether agent is online or offline
- Persistent agents (no support code) stay in session list when disconnected
- Dashboard shows online/offline status with color indicator
- Connect/Chat buttons disabled when agent is offline
- Agent reconnection reuses existing session
---
 server/src/api/mod.rs        |  4 +++
 server/src/relay/mod.rs      |  7 ++--
 server/src/session/mod.rs    | 63 ++++++++++++++++++++++++++++++++++--
 server/static/dashboard.html | 15 ++++++---
 4 files changed, 81 insertions(+), 8 deletions(-)

diff --git a/server/src/api/mod.rs b/server/src/api/mod.rs
index a48e5be..57ddd49 100644
--- a/server/src/api/mod.rs
+++ b/server/src/api/mod.rs
@@ -18,6 +18,8 @@ pub struct SessionInfo {
     pub started_at: String,
     pub viewer_count: usize,
     pub is_streaming: bool,
+    pub is_online: bool,
+    pub is_persistent: bool,
     pub last_heartbeat: String,
     pub os_version: Option<String>,
     pub is_elevated: bool,
@@ -34,6 +36,8 @@ impl From<crate::session::Session> for SessionInfo {
             started_at: s.started_at.to_rfc3339(),
             viewer_count: s.viewer_count,
             is_streaming: s.is_streaming,
+            is_online: s.is_online,
+            is_persistent: s.is_persistent,
             last_heartbeat: s.last_heartbeat.to_rfc3339(),
             os_version: s.os_version,
             is_elevated: s.is_elevated,
diff --git a/server/src/relay/mod.rs b/server/src/relay/mod.rs
index b947220..98e704b 100644
--- a/server/src/relay/mod.rs
+++ b/server/src/relay/mod.rs
@@ -97,7 +97,9 @@ async fn handle_agent_connection(
     }
 
     // Register the agent and get channels
-    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone()).await;
+    // Persistent agents (no support code) keep their session when disconnected
+    let is_persistent = support_code.is_none();
+    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone(), is_persistent).await;
 
     info!("Session created: {} (agent in idle mode)", session_id);
 
@@ -221,7 +223,8 @@ async fn handle_agent_connection(
     // Cleanup
     input_forward.abort();
     cancel_check.abort();
-    sessions_cleanup.remove_session(session_id).await;
+    // Mark agent as disconnected (persistent agents stay in list as offline)
+    sessions_cleanup.mark_agent_disconnected(session_id).await;
 
     // Mark support code as completed if one was used (unless cancelled)
     if let Some(ref code) = support_code_cleanup {
diff --git a/server/src/session/mod.rs b/server/src/session/mod.rs
index 7ceda0a..c6e1fe3 100644
--- a/server/src/session/mod.rs
+++ b/server/src/session/mod.rs
@@ -30,6 +30,8 @@ pub struct Session {
     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>,
@@ -76,7 +78,35 @@ impl SessionManager {
     }
 
     /// Register a new agent and create a session
-    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);
+                    }
+                }
+            }
+        }
+
+        // Create new session
         let session_id = Uuid::new_v4();
 
         // Create channels
@@ -91,6 +121,8 @@ impl SessionManager {
             started_at: now,
             viewer_count: 0,
             is_streaming: false,
+            is_online: true,
+            is_persistent,
             last_heartbeat: now,
             os_version: None,
             is_elevated: false,
@@ -255,10 +287,37 @@ impl SessionManager {
         }
     }
 
-    /// Remove a session (when agent disconnects)
+    /// 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);
         }
diff --git a/server/static/dashboard.html b/server/static/dashboard.html
index 5164955..6f7318c 100644
--- a/server/static/dashboard.html
+++ b/server/static/dashboard.html
@@ -670,12 +670,14 @@
             container.innerHTML = '<div style="padding: 12px;">' + machines.map(m => {
                 const started = new Date(m.started_at).toLocaleString();
                 const isSelected = selectedMachine?.id === m.id;
+                const statusColor = m.is_online ? 'hsl(142, 76%, 50%)' : 'hsl(0, 0%, 50%)';
+                const statusText = m.is_online ? 'Online' : 'Offline';
                 return '<div class="sidebar-item' + (isSelected ? ' active' : '') + '" onclick="selectMachine(\'' + m.id + '\')" style="margin-bottom: 8px; padding: 12px;">' +
                     '<div style="display: flex; align-items: center; gap: 12px;">' +
-                        '<div style="width: 10px; height: 10px; border-radius: 50%; background: hsl(142, 76%, 50%);"></div>' +
+                        '<div style="width: 10px; height: 10px; border-radius: 50%; background: ' + statusColor + ';"></div>' +
                         '<div>' +
                             '<div style="font-weight: 500;">' + (m.agent_name || m.agent_id.slice(0,8)) + '</div>' +
-                            '<div style="font-size: 12px; color: hsl(var(--muted-foreground));">Connected ' + started + '</div>' +
+                            '<div style="font-size: 12px; color: hsl(var(--muted-foreground));">' + statusText + ' • ' + started + '</div>' +
                         '</div>' +
                     '</div>' +
                 '</div>';
@@ -698,10 +700,15 @@
 
             const m = selectedMachine;
             const started = new Date(m.started_at).toLocaleString();
+            const statusColor = m.is_online ? 'hsl(142, 76%, 50%)' : 'hsl(0, 0%, 50%)';
+            const statusText = m.is_online ? 'Online' : 'Offline';
+            const connectDisabled = m.is_online ? '' : 'disabled';
+            const connectTitle = m.is_online ? '' : 'title="Agent is offline"';
 
             container.innerHTML =
                 '<div class="detail-section">' +
                     '<div class="detail-section-title">Machine Info</div>' +
+                    '<div class="detail-row"><span class="detail-label">Status</span><span class="detail-value" style="color: ' + statusColor + ';">' + statusText + '</span></div>' +
                     '<div class="detail-row"><span class="detail-label">Agent ID</span><span class="detail-value">' + m.agent_id.slice(0,8) + '...</span></div>' +
                     '<div class="detail-row"><span class="detail-label">Session ID</span><span class="detail-value">' + m.id.slice(0,8) + '...</span></div>' +
                     '<div class="detail-row"><span class="detail-label">Connected</span><span class="detail-value">' + started + '</span></div>' +
@@ -709,8 +716,8 @@
                 '</div>' +
                 '<div class="detail-section">' +
                     '<div class="detail-section-title">Actions</div>' +
-                    '<button class="btn btn-primary" style="width: 100%; margin-bottom: 8px;" onclick="connectToMachine(\'' + m.id + '\')">Connect</button>' +
-                    '<button class="btn btn-outline" style="width: 100%; margin-bottom: 8px;" onclick="openChat(\'' + m.id + '\', \'' + (m.agent_name || 'Client').replace(/'/g, "\\'") + '\')">Chat</button>' +
+                    '<button class="btn btn-primary" style="width: 100%; margin-bottom: 8px;" onclick="connectToMachine(\'' + m.id + '\')" ' + connectDisabled + ' ' + connectTitle + '>Connect</button>' +
+                    '<button class="btn btn-outline" style="width: 100%; margin-bottom: 8px;" onclick="openChat(\'' + m.id + '\', \'' + (m.agent_name || 'Client').replace(/'/g, "\\'") + '\')" ' + connectDisabled + '>Chat</button>' +
                     '<button class="btn btn-outline" style="width: 100%; margin-bottom: 8px;" disabled>Transfer Files</button>' +
                     '<button class="btn btn-outline" style="width: 100%; color: hsl(0, 62.8%, 50%);" onclick="disconnectMachine(\'' + m.id + '\', \'' + (m.agent_name || m.agent_id).replace(/'/g, "\\'") + '\')">Disconnect</button>' +
                 '</div>';