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

5
agent/src/main.rs
View File

@@ -204,8 +204,9 @@ fn cleanup_on_exit() {
}
async fn run_agent(config: config::Config) -> Result<()> {
// Create session manager
let mut session = session::SessionManager::new(config.clone());
// Create session manager with elevation status
let elevated = is_elevated();
let mut session = session::SessionManager::new(config.clone(), elevated);
let is_support_session = config.support_code.is_some();
let hostname = config.hostname();

390
agent/src/session/mod.rs
View File

@@ -2,8 +2,8 @@
//!
//! Handles the lifecycle of a remote session including:
//! - Connection to server
//! - Authentication
//! - Frame capture and encoding loop
//! - Idle mode (heartbeat only, minimal resources)
//! - Active/streaming mode (capture and send frames)
//! - Input event handling
#[cfg(windows)]
@@ -36,36 +36,58 @@ fn show_debug_console() {
fn show_debug_console() {
// No-op on non-Windows platforms
}
use crate::proto::{Message, message, ChatMessage};
use crate::proto::{Message, message, ChatMessage, AgentStatus, Heartbeat, HeartbeatAck};
use crate::transport::WebSocketTransport;
use crate::tray::{TrayController, TrayAction};
use anyhow::Result;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::mpsc;
// Heartbeat interval (30 seconds)
const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(30);
// Status report interval (60 seconds)
const STATUS_INTERVAL: Duration = Duration::from_secs(60);
/// Session manager handles the remote control session
pub struct SessionManager {
config: Config,
transport: Option<WebSocketTransport>,
state: SessionState,
// Lazy-initialized streaming resources
capturer: Option<Box<dyn Capturer>>,
encoder: Option<Box<dyn Encoder>>,
input: Option<InputController>,
// Streaming state
current_viewer_id: Option<String>,
// System info for status reports
hostname: String,
is_elevated: bool,
start_time: Instant,
}
#[derive(Debug, Clone, PartialEq)]
enum SessionState {
Disconnected,
Connecting,
Connected,
Active,
Idle, // Connected but not streaming - minimal resource usage
Streaming, // Actively capturing and sending frames
}
impl SessionManager {
/// Create a new session manager
pub fn new(config: Config) -> Self {
pub fn new(config: Config, is_elevated: bool) -> Self {
let hostname = config.hostname();
Self {
config,
transport: None,
state: SessionState::Disconnected,
capturer: None,
encoder: None,
input: None,
current_viewer_id: None,
hostname,
is_elevated,
start_time: Instant::now(),
}
}
@@ -73,104 +95,110 @@ impl SessionManager {
pub async fn connect(&mut self) -> Result<()> {
self.state = SessionState::Connecting;
let hostname = self.config.hostname();
let transport = WebSocketTransport::connect(
&self.config.server_url,
&self.config.agent_id,
&self.config.api_key,
Some(&hostname),
Some(&self.hostname),
self.config.support_code.as_deref(),
).await?;
self.transport = Some(transport);
self.state = SessionState::Connected;
self.state = SessionState::Idle; // Start in idle mode
tracing::info!("Connected to server, entering idle mode");
Ok(())
}
/// Run the session main loop
pub async fn run(&mut self) -> Result<()> {
if self.transport.is_none() {
anyhow::bail!("Not connected");
/// Initialize streaming resources (capturer, encoder, input)
fn init_streaming(&mut self) -> Result<()> {
if self.capturer.is_some() {
return Ok(()); // Already initialized
}
self.state = SessionState::Active;
tracing::info!("Initializing streaming resources...");
// Get primary display
let primary_display = capture::primary_display()?;
tracing::info!("Using display: {} ({}x{})", primary_display.name, primary_display.width, primary_display.height);
tracing::info!("Using display: {} ({}x{})",
primary_display.name, primary_display.width, primary_display.height);
// Create capturer
let mut capturer = capture::create_capturer(
let capturer = capture::create_capturer(
primary_display.clone(),
self.config.capture.use_dxgi,
self.config.capture.gdi_fallback,
)?;
self.capturer = Some(capturer);
// Create encoder
let mut encoder = encoder::create_encoder(
let encoder = encoder::create_encoder(
&self.config.encoding.codec,
self.config.encoding.quality,
)?;
self.encoder = Some(encoder);
// Create input controller
let mut input = InputController::new()?;
let input = InputController::new()?;
self.input = Some(input);
// Calculate frame interval
let frame_interval = Duration::from_millis(1000 / self.config.capture.fps as u64);
let mut last_frame_time = Instant::now();
tracing::info!("Streaming resources initialized");
Ok(())
}
// Main loop
loop {
// Check for incoming messages (non-blocking)
// Collect messages first, then release borrow before handling
let messages: Vec<Message> = {
let transport = self.transport.as_mut().unwrap();
let mut msgs = Vec::new();
while let Some(msg) = transport.try_recv()? {
msgs.push(msg);
}
msgs
};
/// Release streaming resources to save CPU/memory when idle
fn release_streaming(&mut self) {
if self.capturer.is_some() {
tracing::info!("Releasing streaming resources");
self.capturer = None;
self.encoder = None;
self.input = None;
self.current_viewer_id = None;
}
}
for msg in messages {
self.handle_message(&mut input, msg)?;
}
/// Get display count for status reports
fn get_display_count(&self) -> i32 {
capture::enumerate_displays().map(|d| d.len() as i32).unwrap_or(1)
}
// Capture and send frame if interval elapsed
if last_frame_time.elapsed() >= frame_interval {
last_frame_time = Instant::now();
/// Send agent status to server
async fn send_status(&mut self) -> Result<()> {
let status = AgentStatus {
hostname: self.hostname.clone(),
os_version: std::env::consts::OS.to_string(),
is_elevated: self.is_elevated,
uptime_secs: self.start_time.elapsed().as_secs() as i64,
display_count: self.get_display_count(),
is_streaming: self.state == SessionState::Streaming,
};
if let Some(frame) = capturer.capture()? {
let encoded = encoder.encode(&frame)?;
let msg = Message {
payload: Some(message::Payload::AgentStatus(status)),
};
// Skip empty frames (no changes)
if encoded.size > 0 {
let msg = Message {
payload: Some(message::Payload::VideoFrame(encoded.frame)),
};
let transport = self.transport.as_mut().unwrap();
transport.send(msg).await?;
}
}
}
// Small sleep to prevent busy loop
tokio::time::sleep(Duration::from_millis(1)).await;
// Check if still connected
if let Some(transport) = self.transport.as_ref() {
if !transport.is_connected() {
tracing::warn!("Connection lost");
break;
}
} else {
tracing::warn!("Transport is None");
break;
}
if let Some(transport) = self.transport.as_mut() {
transport.send(msg).await?;
}
Ok(())
}
/// Send heartbeat to server
async fn send_heartbeat(&mut self) -> Result<()> {
let heartbeat = Heartbeat {
timestamp: chrono::Utc::now().timestamp_millis(),
};
let msg = Message {
payload: Some(message::Payload::Heartbeat(heartbeat)),
};
if let Some(transport) = self.transport.as_mut() {
transport.send(msg).await?;
}
self.state = SessionState::Disconnected;
Ok(())
}
@@ -180,31 +208,14 @@ impl SessionManager {
anyhow::bail!("Not connected");
}
self.state = SessionState::Active;
// Send initial status
self.send_status().await?;
// Get primary display
let primary_display = capture::primary_display()?;
tracing::info!("Using display: {} ({}x{})", primary_display.name, primary_display.width, primary_display.height);
// Create capturer
let mut capturer = capture::create_capturer(
primary_display.clone(),
self.config.capture.use_dxgi,
self.config.capture.gdi_fallback,
)?;
// Create encoder
let mut encoder = encoder::create_encoder(
&self.config.encoding.codec,
self.config.encoding.quality,
)?;
// Create input controller
let mut input = InputController::new()?;
// Calculate frame interval
let frame_interval = Duration::from_millis(1000 / self.config.capture.fps as u64);
// Timing for heartbeat and status
let mut last_heartbeat = Instant::now();
let mut last_status = Instant::now();
let mut last_frame_time = Instant::now();
let frame_interval = Duration::from_millis(1000 / self.config.capture.fps as u64);
// Main loop
loop {
@@ -217,7 +228,6 @@ impl SessionManager {
return Err(anyhow::anyhow!("USER_EXIT: Session ended by user"));
}
TrayAction::ShowDetails => {
// TODO: Show a details dialog
tracing::info!("User requested details (not yet implemented)");
}
TrayAction::ShowDebugWindow => {
@@ -226,14 +236,13 @@ impl SessionManager {
}
}
// Check if exit was requested
if t.exit_requested() {
tracing::info!("Exit requested via tray");
return Err(anyhow::anyhow!("USER_EXIT: Exit requested by user"));
}
}
// Check for incoming messages (non-blocking)
// Process incoming messages
let messages: Vec<Message> = {
let transport = self.transport.as_mut().unwrap();
let mut msgs = Vec::new();
@@ -254,12 +263,56 @@ impl SessionManager {
timestamp: chat_msg.timestamp,
});
}
continue; // Don't pass to handle_message
continue;
}
self.handle_message(&mut input, msg)?;
// Handle control messages that affect state
if let Some(ref payload) = msg.payload {
match payload {
message::Payload::StartStream(start) => {
tracing::info!("StartStream received from viewer: {}", start.viewer_id);
if let Err(e) = self.init_streaming() {
tracing::error!("Failed to init streaming: {}", e);
} else {
self.state = SessionState::Streaming;
self.current_viewer_id = Some(start.viewer_id.clone());
tracing::info!("Now streaming to viewer {}", start.viewer_id);
}
continue;
}
message::Payload::StopStream(stop) => {
tracing::info!("StopStream received for viewer: {}", stop.viewer_id);
// Only stop if it matches current viewer
if self.current_viewer_id.as_ref() == Some(&stop.viewer_id) {
self.release_streaming();
self.state = SessionState::Idle;
tracing::info!("Stopped streaming, returning to idle mode");
}
continue;
}
message::Payload::Heartbeat(hb) => {
// Respond to server heartbeat with ack
let ack = HeartbeatAck {
client_timestamp: hb.timestamp,
server_timestamp: chrono::Utc::now().timestamp_millis(),
};
let ack_msg = Message {
payload: Some(message::Payload::HeartbeatAck(ack)),
};
if let Some(transport) = self.transport.as_mut() {
let _ = transport.send(ack_msg).await;
}
continue;
}
_ => {}
}
}
// Handle other messages (input events, disconnect, etc.)
self.handle_message(msg)?;
}
// Check for outgoing chat messages from user
// Check for outgoing chat messages
if let Some(c) = chat {
if let Some(outgoing) = c.poll_outgoing() {
let chat_proto = ChatMessage {
@@ -276,27 +329,60 @@ impl SessionManager {
}
}
// Capture and send frame if interval elapsed
if last_frame_time.elapsed() >= frame_interval {
last_frame_time = Instant::now();
if let Some(frame) = capturer.capture()? {
let encoded = encoder.encode(&frame)?;
// Skip empty frames (no changes)
if encoded.size > 0 {
let msg = Message {
payload: Some(message::Payload::VideoFrame(encoded.frame)),
};
let transport = self.transport.as_mut().unwrap();
transport.send(msg).await?;
// State-specific behavior
match self.state {
SessionState::Idle => {
// In idle mode, just send heartbeats and status periodically
if last_heartbeat.elapsed() >= HEARTBEAT_INTERVAL {
last_heartbeat = Instant::now();
if let Err(e) = self.send_heartbeat().await {
tracing::warn!("Failed to send heartbeat: {}", e);
}
}
if last_status.elapsed() >= STATUS_INTERVAL {
last_status = Instant::now();
if let Err(e) = self.send_status().await {
tracing::warn!("Failed to send status: {}", e);
}
}
// Longer sleep in idle mode to reduce CPU usage
tokio::time::sleep(Duration::from_millis(100)).await;
}
SessionState::Streaming => {
// In streaming mode, capture and send frames
if last_frame_time.elapsed() >= frame_interval {
last_frame_time = Instant::now();
if let (Some(capturer), Some(encoder)) =
(self.capturer.as_mut(), self.encoder.as_mut())
{
if let Ok(Some(frame)) = capturer.capture() {
if let Ok(encoded) = encoder.encode(&frame) {
if encoded.size > 0 {
let msg = Message {
payload: Some(message::Payload::VideoFrame(encoded.frame)),
};
let transport = self.transport.as_mut().unwrap();
if let Err(e) = transport.send(msg).await {
tracing::warn!("Failed to send frame: {}", e);
}
}
}
}
}
}
// Short sleep in streaming mode
tokio::time::sleep(Duration::from_millis(1)).await;
}
_ => {
// Disconnected or connecting - shouldn't be in main loop
tokio::time::sleep(Duration::from_millis(100)).await;
}
}
// Small sleep to prevent busy loop
tokio::time::sleep(Duration::from_millis(1)).await;
// Check if still connected
if let Some(transport) = self.transport.as_ref() {
if !transport.is_connected() {
@@ -309,70 +395,68 @@ impl SessionManager {
}
}
self.release_streaming();
self.state = SessionState::Disconnected;
Ok(())
}
/// Handle incoming message from server
fn handle_message(&mut self, input: &mut InputController, msg: Message) -> Result<()> {
fn handle_message(&mut self, msg: Message) -> Result<()> {
match msg.payload {
Some(message::Payload::MouseEvent(mouse)) => {
// Handle mouse event
use crate::proto::MouseEventType;
use crate::input::MouseButton;
if let Some(input) = self.input.as_mut() {
use crate::proto::MouseEventType;
use crate::input::MouseButton;
match MouseEventType::try_from(mouse.event_type).unwrap_or(MouseEventType::MouseMove) {
MouseEventType::MouseMove => {
input.mouse_move(mouse.x, mouse.y)?;
}
MouseEventType::MouseDown => {
input.mouse_move(mouse.x, mouse.y)?;
if let Some(ref buttons) = mouse.buttons {
if buttons.left { input.mouse_click(MouseButton::Left, true)?; }
if buttons.right { input.mouse_click(MouseButton::Right, true)?; }
if buttons.middle { input.mouse_click(MouseButton::Middle, true)?; }
match MouseEventType::try_from(mouse.event_type).unwrap_or(MouseEventType::MouseMove) {
MouseEventType::MouseMove => {
input.mouse_move(mouse.x, mouse.y)?;
}
}
MouseEventType::MouseUp => {
if let Some(ref buttons) = mouse.buttons {
if buttons.left { input.mouse_click(MouseButton::Left, false)?; }
if buttons.right { input.mouse_click(MouseButton::Right, false)?; }
if buttons.middle { input.mouse_click(MouseButton::Middle, false)?; }
MouseEventType::MouseDown => {
input.mouse_move(mouse.x, mouse.y)?;
if let Some(ref buttons) = mouse.buttons {
if buttons.left { input.mouse_click(MouseButton::Left, true)?; }
if buttons.right { input.mouse_click(MouseButton::Right, true)?; }
if buttons.middle { input.mouse_click(MouseButton::Middle, true)?; }
}
}
MouseEventType::MouseUp => {
if let Some(ref buttons) = mouse.buttons {
if buttons.left { input.mouse_click(MouseButton::Left, false)?; }
if buttons.right { input.mouse_click(MouseButton::Right, false)?; }
if buttons.middle { input.mouse_click(MouseButton::Middle, false)?; }
}
}
MouseEventType::MouseWheel => {
input.mouse_scroll(mouse.wheel_delta_x, mouse.wheel_delta_y)?;
}
}
MouseEventType::MouseWheel => {
input.mouse_scroll(mouse.wheel_delta_x, mouse.wheel_delta_y)?;
}
}
}
Some(message::Payload::KeyEvent(key)) => {
// Handle keyboard event
input.key_event(key.vk_code as u16, key.down)?;
}
Some(message::Payload::SpecialKey(special)) => {
use crate::proto::SpecialKey;
match SpecialKey::try_from(special.key).ok() {
Some(SpecialKey::CtrlAltDel) => {
input.send_ctrl_alt_del()?;
}
_ => {}
if let Some(input) = self.input.as_mut() {
input.key_event(key.vk_code as u16, key.down)?;
}
}
Some(message::Payload::Heartbeat(_)) => {
// Respond to heartbeat
// TODO: Send heartbeat ack
Some(message::Payload::SpecialKey(special)) => {
if let Some(input) = self.input.as_mut() {
use crate::proto::SpecialKey;
match SpecialKey::try_from(special.key).ok() {
Some(SpecialKey::CtrlAltDel) => {
input.send_ctrl_alt_del()?;
}
_ => {}
}
}
}
Some(message::Payload::Disconnect(disc)) => {
tracing::info!("Disconnect requested: {}", disc.reason);
// Check if this is a cancellation (support session)
if disc.reason.contains("cancelled") {
return Err(anyhow::anyhow!("SESSION_CANCELLED: {}", disc.reason));
}
// Check if this is an admin disconnect (persistent session)
if disc.reason.contains("administrator") || disc.reason.contains("Disconnected") {
return Err(anyhow::anyhow!("ADMIN_DISCONNECT: {}", disc.reason));
}