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

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