azcomputerguru/guru-connect
1
0
Fork 0
Code Issues 4 Pull Requests Actions Packages Projects Releases Wiki Activity

Implement idle/active mode for scalable agent connections

Browse Source 
- Add StartStream/StopStream/AgentStatus messages to protobuf
- Agent now starts in idle mode (heartbeat only, no capture)
- Agent enters streaming mode when viewer connects (StartStream)
- Agent returns to idle when all viewers disconnect (StopStream)
- Server tracks viewer IDs and sends start/stop commands
- Heartbeat mechanism with 90 second timeout detection
- Session API now includes streaming status and agent info

This allows 2000+ agents to connect with minimal bandwidth.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Mike Swanson
2025-12-28 17:24:51 -07:00
parent 5bb5116b92
commit 4417fdfb6e
6 changed files with 455 additions and 176 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
server/src/api/mod.rs
View File

@@ -17,6 +17,12 @@ pub struct SessionInfo {
pub agent_name: String,
pub started_at: String,
pub viewer_count: usize,
pub is_streaming: bool,
pub last_heartbeat: String,
pub os_version: Option<String>,
pub is_elevated: bool,
pub uptime_secs: i64,
pub display_count: i32,
}
impl From<crate::session::Session> for SessionInfo {
@@ -27,6 +33,12 @@ impl From<crate::session::Session> for SessionInfo {
agent_name: s.agent_name,
started_at: s.started_at.to_rfc3339(),
viewer_count: s.viewer_count,
is_streaming: s.is_streaming,
last_heartbeat: s.last_heartbeat.to_rfc3339(),
os_version: s.os_version,
is_elevated: s.is_elevated,
uptime_secs: s.uptime_secs,
display_count: s.display_count,
}
}
}

52
server/src/relay/mod.rs
View File

@@ -14,6 +14,7 @@ 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;
@@ -98,7 +99,7 @@ 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;
info!("Session created: {}", session_id);
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 {
@@ -123,6 +124,7 @@ async fn handle_agent_connection(
});
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();
@@ -154,7 +156,7 @@ async fn handle_agent_connection(
}
});
// Main loop: receive frames from agent and broadcast to viewers
// Main loop: receive messages from agent
while let Some(msg) = ws_receiver.next().await {
match msg {
Ok(Message::Binary(data)) => {
@@ -163,7 +165,7 @@ async fn handle_agent_connection(
Ok(proto_msg) => {
match &proto_msg.payload {
Some(proto::message::Payload::VideoFrame(_)) => {
// Broadcast frame to all viewers
// Broadcast frame to all viewers (only sent when streaming)
let _ = frame_tx.send(data.to_vec());
}
Some(proto::message::Payload::ChatMessage(chat)) => {
@@ -171,6 +173,27 @@ async fn handle_agent_connection(
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;
}
_ => {}
}
}
@@ -226,8 +249,11 @@ async fn handle_viewer_connection(
}
};
// Join the session
let (mut frame_rx, input_tx) = match sessions.join_session(session_id).await {
// 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);
@@ -235,7 +261,7 @@ async fn handle_viewer_connection(
}
};
info!("Viewer joined session: {}", session_id);
info!("Viewer {} joined session: {}", viewer_id, session_id);
let (mut ws_sender, mut ws_receiver) = socket.split();
@@ -249,6 +275,7 @@ async fn handle_viewer_connection(
});
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 {
@@ -259,7 +286,8 @@ async fn handle_viewer_connection(
Ok(proto_msg) => {
match &proto_msg.payload {
Some(proto::message::Payload::MouseEvent(_)) |
Some(proto::message::Payload::KeyEvent(_)) => {
Some(proto::message::Payload::KeyEvent(_)) |
Some(proto::message::Payload::SpecialKey(_)) => {
// Forward input to agent
let _ = input_tx.send(data.to_vec()).await;
}
@@ -277,19 +305,19 @@ async fn handle_viewer_connection(
}
}
Ok(Message::Close(_)) => {
info!("Viewer disconnected from session: {}", session_id);
info!("Viewer {} disconnected from session: {}", viewer_id, session_id);
break;
}
Ok(_) => {}
Err(e) => {
error!("WebSocket error from viewer: {}", e);
error!("WebSocket error from viewer {}: {}", viewer_id, e);
break;
}
}
}
// Cleanup
// Cleanup (this sends StopStream to agent if last viewer)
frame_forward.abort();
sessions_cleanup.leave_session(session_id).await;
info!("Viewer left session: {}", session_id);
sessions_cleanup.leave_session(session_id, &viewer_id_cleanup).await;
info!("Viewer {} left session: {}", viewer_id_cleanup, session_id);
}

148
server/src/session/mod.rs
View File

@@ -3,8 +3,9 @@
//! Manages active remote desktop sessions, tracking which agents
//! are connected and which viewers are watching them.
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::{broadcast, RwLock};
use uuid::Uuid;
@@ -14,6 +15,12 @@ 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 {
@@ -22,6 +29,13 @@ pub struct Session {
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
@@ -40,6 +54,10 @@ struct SessionData {
/// 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
@@ -65,19 +83,28 @@ impl SessionManager {
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: chrono::Utc::now(),
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, // Will be taken by the agent handler
input_rx: None,
viewers: HashSet::new(),
last_heartbeat_instant: Instant::now(),
};
let mut sessions = self.sessions.write().await;
@@ -89,6 +116,59 @@ impl SessionManager {
(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;
@@ -104,24 +184,74 @@ impl SessionManager {
sessions.get(&session_id).map(|s| s.info.clone())
}
/// Join a session as a viewer
pub async fn join_session(&self, session_id: SessionId) -> Option<(FrameReceiver, InputSender)> {
/// 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)?;
session_data.info.viewer_count += 1;
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))
}
/// Leave a session as a viewer
pub async fn leave_session(&self, session_id: SessionId) {
/// 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.info.viewer_count = session_data.info.viewer_count.saturating_sub(1);
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;
}
}