     1→//! GuruConnect Agent - Remote Desktop Agent for Windows
     2→//!
     3→//! Provides screen capture, input injection, and remote control capabilities.
     4→//!
     5→//! Usage:
     6→//!   guruconnect-agent.exe [support_code]
     7→//!
     8→//! If a support code is provided, the agent will connect using that code
     9→//! for a one-time support session.
    10→
    11→// Hide console window by default on Windows (release builds)
    12→#![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]
    13→
    14→mod capture;
    15→mod chat;
    16→mod config;
    17→mod encoder;
    18→mod input;
    19→mod session;
    20→mod startup;
    21→mod transport;
    22→mod tray;
    23→
    24→pub mod proto {
    25→    include!(concat!(env!("OUT_DIR"), "/guruconnect.rs"));
    26→}
    27→
    28→use anyhow::Result;
    29→use tracing::{info, error, warn, Level};
    30→use tracing_subscriber::FmtSubscriber;
    31→
    32→#[cfg(windows)]
    33→use windows::Win32::UI::WindowsAndMessaging::{MessageBoxW, MB_OK, MB_ICONINFORMATION};
    34→#[cfg(windows)]
    35→use windows::core::PCWSTR;
    36→#[cfg(windows)]
    37→use windows::Win32::Security::{GetTokenInformation, TokenElevation, TOKEN_ELEVATION, TOKEN_QUERY};
    38→#[cfg(windows)]
    39→use windows::Win32::System::Threading::{GetCurrentProcess, OpenProcessToken};
    40→#[cfg(windows)]
    41→use windows::Win32::System::Console::{AllocConsole, FreeConsole, GetConsoleWindow};
    42→#[cfg(windows)]
    43→use windows::Win32::UI::WindowsAndMessaging::{ShowWindow, SW_SHOW, SW_HIDE};
    44→
    45→/// Extract a 6-digit support code from the executable's filename.
    46→/// Looks for patterns like "GuruConnect-123456.exe" or "123456.exe"
    47→fn extract_code_from_filename() -> Option<String> {
    48→    // Get the path to the current executable
    49→    let exe_path = std::env::current_exe().ok()?;
    50→    let filename = exe_path.file_stem()?.to_str()?;
    51→
    52→    // Look for a 6-digit number in the filename
    53→    // Try common patterns: "Name-123456", "Name_123456", "123456"
    54→    for part in filename.split(|c| c == '-' || c == '_' || c == '.') {
    55→        let trimmed = part.trim();
    56→        if trimmed.len() == 6 && trimmed.chars().all(|c| c.is_ascii_digit()) {
    57→            return Some(trimmed.to_string());
    58→        }
    59→    }
    60→
    61→    // Also check if the last 6 characters are digits (e.g., "GuruConnect123456")
    62→    if filename.len() >= 6 {
    63→        let last_six = &filename[filename.len() - 6..];
    64→        if last_six.chars().all(|c| c.is_ascii_digit()) {
    65→            return Some(last_six.to_string());
    66→        }
    67→    }
    68→
    69→    None
    70→}
    71→
    72→/// Check if the process is running with elevated privileges (Windows only)
    73→#[cfg(windows)]
    74→fn is_elevated() -> bool {
    75→    unsafe {
    76→        let mut token_handle = windows::Win32::Foundation::HANDLE::default();
    77→        if OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &mut token_handle).is_err() {
    78→            return false;
    79→        }
    80→
    81→        let mut elevation = TOKEN_ELEVATION::default();
    82→        let mut size = std::mem::size_of::<TOKEN_ELEVATION>() as u32;
    83→
    84→        let result = GetTokenInformation(
    85→            token_handle,
    86→            TokenElevation,
    87→            Some(&mut elevation as *mut _ as *mut _),
    88→            size,
    89→            &mut size,
    90→        );
    91→
    92→        let _ = windows::Win32::Foundation::CloseHandle(token_handle);
    93→
    94→        result.is_ok() && elevation.TokenIsElevated != 0
    95→    }
    96→}
    97→
    98→#[cfg(not(windows))]
    99→fn is_elevated() -> bool {
   100→    // On non-Windows, check if running as root
   101→    unsafe { libc::geteuid() == 0 }
   102→}
   103→
   104→/// Show a message box to the user (Windows only)
   105→#[cfg(windows)]
   106→fn show_message_box(title: &str, message: &str) {
   107→    use std::ffi::OsStr;
   108→    use std::os::windows::ffi::OsStrExt;
   109→
   110→    // Convert strings to wide strings for Windows API
   111→    let title_wide: Vec<u16> = OsStr::new(title)
   112→        .encode_wide()
   113→        .chain(std::iter::once(0))
   114→        .collect();
   115→    let message_wide: Vec<u16> = OsStr::new(message)
   116→        .encode_wide()
   117→        .chain(std::iter::once(0))
   118→        .collect();
   119→
   120→    unsafe {
   121→        MessageBoxW(
   122→            None,
   123→            PCWSTR(message_wide.as_ptr()),
   124→            PCWSTR(title_wide.as_ptr()),
   125→            MB_OK | MB_ICONINFORMATION,
   126→        );
   127→    }
   128→}
   129→
   130→#[cfg(not(windows))]
   131→fn show_message_box(_title: &str, _message: &str) {
   132→    // No-op on non-Windows platforms
   133→}
   134→
   135→/// Show the debug console window (Windows only)
   136→#[cfg(windows)]
   137→fn show_debug_console() {
   138→    unsafe {
   139→        // Check if we already have a console
   140→        let hwnd = GetConsoleWindow();
   141→        if hwnd.0 == std::ptr::null_mut() {
   142→            // No console, allocate one
   143→            let _ = AllocConsole();
   144→            info!("Debug console window opened");
   145→        } else {
   146→            // Console exists, make sure it's visible
   147→            let _ = ShowWindow(hwnd, SW_SHOW);
   148→            info!("Debug console window shown");
   149→        }
   150→    }
   151→}
   152→
   153→#[cfg(not(windows))]
   154→fn show_debug_console() {
   155→    // No-op on non-Windows platforms
   156→}
   157→
   158→#[tokio::main]
   159→async fn main() -> Result<()> {
   160→    // Initialize logging
   161→    let subscriber = FmtSubscriber::builder()
   162→        .with_max_level(Level::INFO)
   163→        .with_target(true)
   164→        .with_thread_ids(true)
   165→        .init();
   166→
   167→    info!("GuruConnect Agent v{}", env!("CARGO_PKG_VERSION"));
   168→
   169→    // Check and log elevation status
   170→    if is_elevated() {
   171→        info!("Running with elevated (administrator) privileges");
   172→    } else {
   173→        info!("Running with standard user privileges");
   174→    }
   175→
   176→    // Extract support code from executable filename
   177→    // e.g., GuruConnect-123456.exe -> 123456
   178→    let support_code = extract_code_from_filename();
   179→    if let Some(ref code) = support_code {
   180→        info!("Support code from filename: {}", code);
   181→    }
   182→
   183→    // Load configuration
   184→    let mut config = config::Config::load()?;
   185→    config.support_code = support_code;
   186→    info!("Loaded configuration for server: {}", config.server_url);
   187→
   188→    // Run the agent
   189→    if let Err(e) = run_agent(config).await {
   190→        error!("Agent error: {}", e);
   191→        return Err(e);
   192→    }
   193→
   194→    Ok(())
   195→}
   196→
   197→/// Clean up before exiting (remove from startup, etc.)
   198→/// Called when user explicitly ends session or support session completes
   199→fn cleanup_on_exit() {
   200→    info!("Cleaning up before exit");
   201→    if let Err(e) = startup::remove_from_startup() {
   202→        warn!("Failed to remove from startup: {}", e);
   203→    }
   204→}
   205→
   206→async fn run_agent(config: config::Config) -> Result<()> {
   207→    // Create session manager
   208→    let mut session = session::SessionManager::new(config.clone());
   209→    let is_support_session = config.support_code.is_some();
   210→    let hostname = config.hostname();
   211→
   212→    // Add to startup so we reconnect after reboot
   213→    // Persistent agents (no support code) should ALWAYS be in startup
   214→    // Support sessions only need startup temporarily while active
   215→    if let Err(e) = startup::add_to_startup() {
   216→        warn!("Failed to add to startup: {}. Agent won't persist through reboot.", e);
   217→    }
   218→
   219→    // Create tray icon
   220→    // Only support sessions can be ended by user - persistent agents are admin-managed
   221→    let tray = match tray::TrayController::new(&hostname, config.support_code.as_deref(), is_support_session) {
   222→        Ok(t) => {
   223→            info!("Tray icon created");
   224→            Some(t)
   225→        }
   226→        Err(e) => {
   227→            warn!("Failed to create tray icon: {}. Continuing without tray.", e);
   228→            None
   229→        }
   230→    };
   231→
   232→    // Create chat controller
   233→    let chat_ctrl = chat::ChatController::new();
   234→    if chat_ctrl.is_some() {
   235→        info!("Chat controller created");
   236→    }
   237→
   238→    // Connect to server and run main loop
   239→    loop {
   240→        info!("Connecting to server...");
   241→
   242→        // Check if user requested exit via tray before connecting (support sessions only)
   243→        if is_support_session {
   244→            if let Some(ref t) = tray {
   245→                if t.exit_requested() {
   246→                    info!("Exit requested by user");
   247→                    cleanup_on_exit();
   248→                    return Ok(());
   249→                }
   250→            }
   251→        }
   252→
   253→        match session.connect().await {
   254→            Ok(_) => {
   255→                info!("Connected to server");
   256→
   257→                // Update tray status
   258→                if let Some(ref t) = tray {
   259→                    t.update_status("Status: Connected");
   260→                }
   261→
   262→                // Run session until disconnect, passing tray and chat for event processing
   263→                if let Err(e) = session.run_with_tray(tray.as_ref(), chat_ctrl.as_ref()).await {
   264→                    let error_msg = e.to_string();
   265→
   266→                    // Check if this is a user-initiated exit
   267→                    if error_msg.contains("USER_EXIT") {
   268→                        info!("Session ended by user");
   269→                        cleanup_on_exit();
   270→                        return Ok(());
   271→                    }
   272→
   273→                    // Check if this is a cancellation
   274→                    if error_msg.contains("SESSION_CANCELLED") {
   275→                        info!("Session was cancelled by technician");
   276→                        cleanup_on_exit();
   277→                        show_message_box(
   278→                            "Support Session Ended",
   279→                            "The support session was cancelled by the technician.\n\nThis window will close automatically.",
   280→                        );
   281→                        // Exit cleanly without reconnecting
   282→                        return Ok(());
   283→                    }
   284→
   285→                    // Check if this is an admin disconnect (uninstall)
   286→                    if error_msg.contains("ADMIN_DISCONNECT") {
   287→                        info!("Session was disconnected by administrator - uninstalling");
   288→                        if let Err(e) = startup::uninstall() {
   289→                            warn!("Uninstall failed: {}", e);
   290→                        }
   291→                        show_message_box(
   292→                            "Remote Session Ended",
   293→                            "The remote support session has been ended by the administrator.\n\nThe agent will be removed from this computer.",
   294→                        );
   295→                        return Ok(());
   296→                    }
   297→
   298→                    error!("Session error: {}", e);
   299→                }
   300→            }
   301→            Err(e) => {
   302→                let error_msg = e.to_string();
   303→
   304→                // Check if connection was rejected due to cancelled code
   305→                if error_msg.contains("cancelled") {
   306→                    info!("Support code was cancelled before connection");
   307→                    cleanup_on_exit();
   308→                    show_message_box(
   309→                        "Support Session Cancelled",
   310→                        "This support session has been cancelled.\n\nPlease contact your technician for a new support code.",
   311→                    );
   312→                    // Exit cleanly without reconnecting
   313→                    return Ok(());
   314→                }
   315→
   316→                error!("Connection failed: {}", e);
   317→            }
   318→        }
   319→
   320→        // For support sessions, don't reconnect if something goes wrong
   321→        if is_support_session {
   322→            info!("Support session ended, not reconnecting");
   323→            cleanup_on_exit();
   324→            return Ok(());
   325→        }
   326→
   327→        // Wait before reconnecting (persistent agents only - support sessions already exited above)
   328→        info!("Reconnecting in 5 seconds...");
   329→        tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
   330→    }
   331→}
   332→

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

//! GuruConnect Agent - Remote Desktop Agent for Windows
//!
//! Provides screen capture, input injection, and remote control capabilities.
//!
//! Usage:
//!   guruconnect-agent.exe [support_code]
//!
//! If a support code is provided, the agent will connect using that code
//! for a one-time support session.

// Hide console window by default on Windows (release builds)
#![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]

mod capture;
mod chat;
mod config;
mod encoder;
mod input;
mod session;
mod startup;
mod transport;
mod tray;

pub mod proto {
    include!(concat!(env!("OUT_DIR"), "/guruconnect.rs"));
}

use anyhow::Result;
use tracing::{info, error, warn, Level};
use tracing_subscriber::FmtSubscriber;

#[cfg(windows)]
use windows::Win32::UI::WindowsAndMessaging::{MessageBoxW, MB_OK, MB_ICONINFORMATION};
#[cfg(windows)]
use windows::core::PCWSTR;
#[cfg(windows)]
use windows::Win32::Security::{GetTokenInformation, TokenElevation, TOKEN_ELEVATION, TOKEN_QUERY};
#[cfg(windows)]
use windows::Win32::System::Threading::{GetCurrentProcess, OpenProcessToken};
#[cfg(windows)]
use windows::Win32::System::Console::{AllocConsole, FreeConsole, GetConsoleWindow};
#[cfg(windows)]
use windows::Win32::UI::WindowsAndMessaging::{ShowWindow, SW_SHOW, SW_HIDE};

/// Extract a 6-digit support code from the executable's filename.
/// Looks for patterns like "GuruConnect-123456.exe" or "123456.exe"
fn extract_code_from_filename() -> Option<String> {
    // Get the path to the current executable
    let exe_path = std::env::current_exe().ok()?;
    let filename = exe_path.file_stem()?.to_str()?;

    // Look for a 6-digit number in the filename
    // Try common patterns: "Name-123456", "Name_123456", "123456"
    for part in filename.split(|c| c == '-' || c == '_' || c == '.') {
        let trimmed = part.trim();
        if trimmed.len() == 6 && trimmed.chars().all(|c| c.is_ascii_digit()) {
            return Some(trimmed.to_string());
        }
    }

    // Also check if the last 6 characters are digits (e.g., "GuruConnect123456")
    if filename.len() >= 6 {
        let last_six = &filename[filename.len() - 6..];
        if last_six.chars().all(|c| c.is_ascii_digit()) {
            return Some(last_six.to_string());
        }
    }

    None
}

/// Check if the process is running with elevated privileges (Windows only)
#[cfg(windows)]
fn is_elevated() -> bool {
    unsafe {
        let mut token_handle = windows::Win32::Foundation::HANDLE::default();
        if OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &mut token_handle).is_err() {
            return false;
        }

        let mut elevation = TOKEN_ELEVATION::default();
        let mut size = std::mem::size_of::<TOKEN_ELEVATION>() as u32;

        let result = GetTokenInformation(
            token_handle,
            TokenElevation,
            Some(&mut elevation as *mut _ as *mut _),
            size,
            &mut size,
        );

        let _ = windows::Win32::Foundation::CloseHandle(token_handle);

        result.is_ok() && elevation.TokenIsElevated != 0
    }
}

#[cfg(not(windows))]
fn is_elevated() -> bool {
    // On non-Windows, check if running as root
    unsafe { libc::geteuid() == 0 }
}

/// Show a message box to the user (Windows only)
#[cfg(windows)]
fn show_message_box(title: &str, message: &str) {
    use std::ffi::OsStr;
    use std::os::windows::ffi::OsStrExt;

    // Convert strings to wide strings for Windows API
    let title_wide: Vec<u16> = OsStr::new(title)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();
    let message_wide: Vec<u16> = OsStr::new(message)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    unsafe {
        MessageBoxW(
            None,
            PCWSTR(message_wide.as_ptr()),
            PCWSTR(title_wide.as_ptr()),
            MB_OK | MB_ICONINFORMATION,
        );
    }
}

#[cfg(not(windows))]
fn show_message_box(_title: &str, _message: &str) {
    // No-op on non-Windows platforms
}

/// Show the debug console window (Windows only)
#[cfg(windows)]
fn show_debug_console() {
    unsafe {
        // Check if we already have a console
        let hwnd = GetConsoleWindow();
        if hwnd.0 == std::ptr::null_mut() {
            // No console, allocate one
            let _ = AllocConsole();
            info!("Debug console window opened");
        } else {
            // Console exists, make sure it's visible
            let _ = ShowWindow(hwnd, SW_SHOW);
            info!("Debug console window shown");
        }
    }
}

#[cfg(not(windows))]
fn show_debug_console() {
    // No-op on non-Windows platforms
}

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize logging
    let subscriber = FmtSubscriber::builder()
        .with_max_level(Level::INFO)
        .with_target(true)
        .with_thread_ids(true)
        .init();

    info!("GuruConnect Agent v{}", env!("CARGO_PKG_VERSION"));

    // Check and log elevation status
    if is_elevated() {
        info!("Running with elevated (administrator) privileges");
    } else {
        info!("Running with standard user privileges");
    }

    // Extract support code from executable filename
    // e.g., GuruConnect-123456.exe -> 123456
    let support_code = extract_code_from_filename();
    if let Some(ref code) = support_code {
        info!("Support code from filename: {}", code);
    }

    // Load configuration
    let mut config = config::Config::load()?;
    config.support_code = support_code;
    info!("Loaded configuration for server: {}", config.server_url);

    // Run the agent
    if let Err(e) = run_agent(config).await {
        error!("Agent error: {}", e);
        return Err(e);
    }

    Ok(())
}

/// Clean up before exiting (remove from startup, etc.)
/// Called when user explicitly ends session or support session completes
fn cleanup_on_exit() {
    info!("Cleaning up before exit");
    if let Err(e) = startup::remove_from_startup() {
        warn!("Failed to remove from startup: {}", e);
    }
}

async fn run_agent(config: config::Config) -> Result<()> {
    // Create session manager
    let mut session = session::SessionManager::new(config.clone());
    let is_support_session = config.support_code.is_some();
    let hostname = config.hostname();

    // Add to startup so we reconnect after reboot
    // Persistent agents (no support code) should ALWAYS be in startup
    // Support sessions only need startup temporarily while active
    if let Err(e) = startup::add_to_startup() {
        warn!("Failed to add to startup: {}. Agent won't persist through reboot.", e);
    }

    // Create tray icon
    // Only support sessions can be ended by user - persistent agents are admin-managed
    let tray = match tray::TrayController::new(&hostname, config.support_code.as_deref(), is_support_session) {
        Ok(t) => {
            info!("Tray icon created");
            Some(t)
        }
        Err(e) => {
            warn!("Failed to create tray icon: {}. Continuing without tray.", e);
            None
        }
    };

    // Create chat controller
    let chat_ctrl = chat::ChatController::new();
    if chat_ctrl.is_some() {
        info!("Chat controller created");
    }

    // Connect to server and run main loop
    loop {
        info!("Connecting to server...");

        // Check if user requested exit via tray before connecting (support sessions only)
        if is_support_session {
            if let Some(ref t) = tray {
                if t.exit_requested() {
                    info!("Exit requested by user");
                    cleanup_on_exit();
                    return Ok(());
                }
            }
        }

        match session.connect().await {
            Ok(_) => {
                info!("Connected to server");

                // Update tray status
                if let Some(ref t) = tray {
                    t.update_status("Status: Connected");
                }

                // Run session until disconnect, passing tray and chat for event processing
                if let Err(e) = session.run_with_tray(tray.as_ref(), chat_ctrl.as_ref()).await {
                    let error_msg = e.to_string();

                    // Check if this is a user-initiated exit
                    if error_msg.contains("USER_EXIT") {
                        info!("Session ended by user");
                        cleanup_on_exit();
                        return Ok(());
                    }

                    // Check if this is a cancellation
                    if error_msg.contains("SESSION_CANCELLED") {
                        info!("Session was cancelled by technician");
                        cleanup_on_exit();
                        show_message_box(
                            "Support Session Ended",
                            "The support session was cancelled by the technician.\n\nThis window will close automatically.",
                        );
                        // Exit cleanly without reconnecting
                        return Ok(());
                    }

                    // Check if this is an admin disconnect (uninstall)
                    if error_msg.contains("ADMIN_DISCONNECT") {
                        info!("Session was disconnected by administrator - uninstalling");
                        if let Err(e) = startup::uninstall() {
                            warn!("Uninstall failed: {}", e);
                        }
                        show_message_box(
                            "Remote Session Ended",
                            "The remote support session has been ended by the administrator.\n\nThe agent will be removed from this computer.",
                        );
                        return Ok(());
                    }

                    error!("Session error: {}", e);
                }
            }
            Err(e) => {
                let error_msg = e.to_string();

                // Check if connection was rejected due to cancelled code
                if error_msg.contains("cancelled") {
                    info!("Support code was cancelled before connection");
                    cleanup_on_exit();
                    show_message_box(
                        "Support Session Cancelled",
                        "This support session has been cancelled.\n\nPlease contact your technician for a new support code.",
                    );
                    // Exit cleanly without reconnecting
                    return Ok(());
                }

                error!("Connection failed: {}", e);
            }
        }

        // For support sessions, don't reconnect if something goes wrong
        if is_support_session {
            info!("Support session ended, not reconnecting");
            cleanup_on_exit();
            return Ok(());
        }

        // Wait before reconnecting (persistent agents only - support sessions already exited above)
        info!("Reconnecting in 5 seconds...");
        tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
    }
}


I need to update the main.rs to pass `is_elevated()` to `SessionManager::new()`. The change is on line 208.


Edit

async fn run_agent(config: config::Config) -> Result<()> {
    // Create session manager
    let mut session = session::SessionManager::new(config.clone());
    let is_support_session = config.support_code.is_some();

async fn run_agent(config: config::Config) -> Result<()> {
    // 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();


The file C:\Users\MikeSwanson\Claude\guru-connect\agent\src\main.rs has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
   202→        warn!("Failed to remove from startup: {}", e);
   203→    }
   204→}
   205→
   206→async fn run_agent(config: config::Config) -> Result<()> {
   207→    // Create session manager with elevation status
   208→    let elevated = is_elevated();
   209→    let mut session = session::SessionManager::new(config.clone(), elevated);
   210→    let is_support_session = config.support_code.is_some();
   211→    let hostname = config.hostname();
   212→
   213→    // Add to startup so we reconnect after reboot
   214→    // Persistent agents (no support code) should ALWAYS be in startup

async fn run_agent(config: config::Config) -> Result<()> {
    // 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();

//! GuruConnect Agent - Remote Desktop Agent for Windows
//!
//! Provides screen capture, input injection, and remote control capabilities.
//!
//! Usage:
//!   guruconnect-agent.exe [support_code]
//!
//! If a support code is provided, the agent will connect using that code
//! for a one-time support session.

// Hide console window by default on Windows (release builds)
#![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]

mod capture;
mod chat;
mod config;
mod encoder;
mod input;
mod session;
mod startup;
mod transport;
mod tray;

pub mod proto {
    include!(concat!(env!("OUT_DIR"), "/guruconnect.rs"));
}

use anyhow::Result;
use tracing::{info, error, warn, Level};
use tracing_subscriber::FmtSubscriber;

#[cfg(windows)]
use windows::Win32::UI::WindowsAndMessaging::{MessageBoxW, MB_OK, MB_ICONINFORMATION};
#[cfg(windows)]
use windows::core::PCWSTR;
#[cfg(windows)]
use windows::Win32::Security::{GetTokenInformation, TokenElevation, TOKEN_ELEVATION, TOKEN_QUERY};
#[cfg(windows)]
use windows::Win32::System::Threading::{GetCurrentProcess, OpenProcessToken};
#[cfg(windows)]
use windows::Win32::System::Console::{AllocConsole, FreeConsole, GetConsoleWindow};
#[cfg(windows)]
use windows::Win32::UI::WindowsAndMessaging::{ShowWindow, SW_SHOW, SW_HIDE};

/// Extract a 6-digit support code from the executable's filename.
/// Looks for patterns like "GuruConnect-123456.exe" or "123456.exe"
fn extract_code_from_filename() -> Option<String> {
    // Get the path to the current executable
    let exe_path = std::env::current_exe().ok()?;
    let filename = exe_path.file_stem()?.to_str()?;

    // Look for a 6-digit number in the filename
    // Try common patterns: "Name-123456", "Name_123456", "123456"
    for part in filename.split(|c| c == '-' || c == '_' || c == '.') {
        let trimmed = part.trim();
        if trimmed.len() == 6 && trimmed.chars().all(|c| c.is_ascii_digit()) {
            return Some(trimmed.to_string());
        }
    }

    // Also check if the last 6 characters are digits (e.g., "GuruConnect123456")
    if filename.len() >= 6 {
        let last_six = &filename[filename.len() - 6..];
        if last_six.chars().all(|c| c.is_ascii_digit()) {
            return Some(last_six.to_string());
        }
    }

    None
}

/// Check if the process is running with elevated privileges (Windows only)
#[cfg(windows)]
fn is_elevated() -> bool {
    unsafe {
        let mut token_handle = windows::Win32::Foundation::HANDLE::default();
        if OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &mut token_handle).is_err() {
            return false;
        }

        let mut elevation = TOKEN_ELEVATION::default();
        let mut size = std::mem::size_of::<TOKEN_ELEVATION>() as u32;

        let result = GetTokenInformation(
            token_handle,
            TokenElevation,
            Some(&mut elevation as *mut _ as *mut _),
            size,
            &mut size,
        );

        let _ = windows::Win32::Foundation::CloseHandle(token_handle);

        result.is_ok() && elevation.TokenIsElevated != 0
    }
}

#[cfg(not(windows))]
fn is_elevated() -> bool {
    // On non-Windows, check if running as root
    unsafe { libc::geteuid() == 0 }
}

/// Show a message box to the user (Windows only)
#[cfg(windows)]
fn show_message_box(title: &str, message: &str) {
    use std::ffi::OsStr;
    use std::os::windows::ffi::OsStrExt;

    // Convert strings to wide strings for Windows API
    let title_wide: Vec<u16> = OsStr::new(title)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();
    let message_wide: Vec<u16> = OsStr::new(message)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    unsafe {
        MessageBoxW(
            None,
            PCWSTR(message_wide.as_ptr()),
            PCWSTR(title_wide.as_ptr()),
            MB_OK | MB_ICONINFORMATION,
        );
    }
}

#[cfg(not(windows))]
fn show_message_box(_title: &str, _message: &str) {
    // No-op on non-Windows platforms
}

/// Show the debug console window (Windows only)
#[cfg(windows)]
fn show_debug_console() {
    unsafe {
        // Check if we already have a console
        let hwnd = GetConsoleWindow();
        if hwnd.0 == std::ptr::null_mut() {
            // No console, allocate one
            let _ = AllocConsole();
            info!("Debug console window opened");
        } else {
            // Console exists, make sure it's visible
            let _ = ShowWindow(hwnd, SW_SHOW);
            info!("Debug console window shown");
        }
    }
}

#[cfg(not(windows))]
fn show_debug_console() {
    // No-op on non-Windows platforms
}

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize logging
    let subscriber = FmtSubscriber::builder()
        .with_max_level(Level::INFO)
        .with_target(true)
        .with_thread_ids(true)
        .init();

    info!("GuruConnect Agent v{}", env!("CARGO_PKG_VERSION"));

    // Check and log elevation status
    if is_elevated() {
        info!("Running with elevated (administrator) privileges");
    } else {
        info!("Running with standard user privileges");
    }

    // Extract support code from executable filename
    // e.g., GuruConnect-123456.exe -> 123456
    let support_code = extract_code_from_filename();
    if let Some(ref code) = support_code {
        info!("Support code from filename: {}", code);
    }

    // Load configuration
    let mut config = config::Config::load()?;
    config.support_code = support_code;
    info!("Loaded configuration for server: {}", config.server_url);

    // Run the agent
    if let Err(e) = run_agent(config).await {
        error!("Agent error: {}", e);
        return Err(e);
    }

    Ok(())
}

/// Clean up before exiting (remove from startup, etc.)
/// Called when user explicitly ends session or support session completes
fn cleanup_on_exit() {
    info!("Cleaning up before exit");
    if let Err(e) = startup::remove_from_startup() {
        warn!("Failed to remove from startup: {}", e);
    }
}

async fn run_agent(config: config::Config) -> Result<()> {
    // Create session manager
    let mut session = session::SessionManager::new(config.clone());
    let is_support_session = config.support_code.is_some();
    let hostname = config.hostname();

    // Add to startup so we reconnect after reboot
    // Persistent agents (no support code) should ALWAYS be in startup
    // Support sessions only need startup temporarily while active
    if let Err(e) = startup::add_to_startup() {
        warn!("Failed to add to startup: {}. Agent won't persist through reboot.", e);
    }

    // Create tray icon
    // Only support sessions can be ended by user - persistent agents are admin-managed
    let tray = match tray::TrayController::new(&hostname, config.support_code.as_deref(), is_support_session) {
        Ok(t) => {
            info!("Tray icon created");
            Some(t)
        }
        Err(e) => {
            warn!("Failed to create tray icon: {}. Continuing without tray.", e);
            None
        }
    };

    // Create chat controller
    let chat_ctrl = chat::ChatController::new();
    if chat_ctrl.is_some() {
        info!("Chat controller created");
    }

    // Connect to server and run main loop
    loop {
        info!("Connecting to server...");

        // Check if user requested exit via tray before connecting (support sessions only)
        if is_support_session {
            if let Some(ref t) = tray {
                if t.exit_requested() {
                    info!("Exit requested by user");
                    cleanup_on_exit();
                    return Ok(());
                }
            }
        }

        match session.connect().await {
            Ok(_) => {
                info!("Connected to server");

                // Update tray status
                if let Some(ref t) = tray {
                    t.update_status("Status: Connected");
                }

                // Run session until disconnect, passing tray and chat for event processing
                if let Err(e) = session.run_with_tray(tray.as_ref(), chat_ctrl.as_ref()).await {
                    let error_msg = e.to_string();

                    // Check if this is a user-initiated exit
                    if error_msg.contains("USER_EXIT") {
                        info!("Session ended by user");
                        cleanup_on_exit();
                        return Ok(());
                    }

                    // Check if this is a cancellation
                    if error_msg.contains("SESSION_CANCELLED") {
                        info!("Session was cancelled by technician");
                        cleanup_on_exit();
                        show_message_box(
                            "Support Session Ended",
                            "The support session was cancelled by the technician.\n\nThis window will close automatically.",
                        );
                        // Exit cleanly without reconnecting
                        return Ok(());
                    }

                    // Check if this is an admin disconnect (uninstall)
                    if error_msg.contains("ADMIN_DISCONNECT") {
                        info!("Session was disconnected by administrator - uninstalling");
                        if let Err(e) = startup::uninstall() {
                            warn!("Uninstall failed: {}", e);
                        }
                        show_message_box(
                            "Remote Session Ended",
                            "The remote support session has been ended by the administrator.\n\nThe agent will be removed from this computer.",
                        );
                        return Ok(());
                    }

                    error!("Session error: {}", e);
                }
            }
            Err(e) => {
                let error_msg = e.to_string();

                // Check if connection was rejected due to cancelled code
                if error_msg.contains("cancelled") {
                    info!("Support code was cancelled before connection");
                    cleanup_on_exit();
                    show_message_box(
                        "Support Session Cancelled",
                        "This support session has been cancelled.\n\nPlease contact your technician for a new support code.",
                    );
                    // Exit cleanly without reconnecting
                    return Ok(());
                }

                error!("Connection failed: {}", e);
            }
        }

        // For support sessions, don't reconnect if something goes wrong
        if is_support_session {
            info!("Support session ended, not reconnecting");
            cleanup_on_exit();
            return Ok(());
        }

        // Wait before reconnecting (persistent agents only - support sessions already exited above)
        info!("Reconnecting in 5 seconds...");
        tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
    }
}

-    let mut session = session::SessionManager::new(config.clone());

+    // Create session manager with elevation status

+    let mut session = session::SessionManager::new(config.clone(), elevated);


Now I need to update the server to:
1. Store agent connections with their status
2. Handle AgentStatus messages
3. Send StartStream when viewer connects
4. Send StopStream when viewer disconnects
5. Implement heartbeat timeout detection

Let me look at the server's session manager and relay modules.


Now let me update the server to handle the new messages and routing. First, let me check the server's session and relay modules:


     1→//! Session management for GuruConnect
     2→//!
     3→//! Manages active remote desktop sessions, tracking which agents
     4→//! are connected and which viewers are watching them.
     5→
     6→use std::collections::HashMap;
     7→use std::sync::Arc;
     8→use tokio::sync::{broadcast, RwLock};
     9→use uuid::Uuid;
    10→
    11→/// Unique identifier for a session
    12→pub type SessionId = Uuid;
    13→
    14→/// Unique identifier for an agent
    15→pub type AgentId = String;
    16→
    17→/// Session state
    18→#[derive(Debug, Clone)]
    19→pub struct Session {
    20→    pub id: SessionId,
    21→    pub agent_id: AgentId,
    22→    pub agent_name: String,
    23→    pub started_at: chrono::DateTime<chrono::Utc>,
    24→    pub viewer_count: usize,
    25→}
    26→
    27→/// Channel for sending frames from agent to viewers
    28→pub type FrameSender = broadcast::Sender<Vec<u8>>;
    29→pub type FrameReceiver = broadcast::Receiver<Vec<u8>>;
    30→
    31→/// Channel for sending input events from viewer to agent
    32→pub type InputSender = tokio::sync::mpsc::Sender<Vec<u8>>;
    33→pub type InputReceiver = tokio::sync::mpsc::Receiver<Vec<u8>>;
    34→
    35→/// Internal session data with channels
    36→struct SessionData {
    37→    info: Session,
    38→    /// Channel for video frames (agent -> viewers)
    39→    frame_tx: FrameSender,
    40→    /// Channel for input events (viewer -> agent)
    41→    input_tx: InputSender,
    42→    input_rx: Option<InputReceiver>,
    43→}
    44→
    45→/// Manages all active sessions
    46→#[derive(Clone)]
    47→pub struct SessionManager {
    48→    sessions: Arc<RwLock<HashMap<SessionId, SessionData>>>,
    49→    agents: Arc<RwLock<HashMap<AgentId, SessionId>>>,
    50→}
    51→
    52→impl SessionManager {
    53→    pub fn new() -> Self {
    54→        Self {
    55→            sessions: Arc::new(RwLock::new(HashMap::new())),
    56→            agents: Arc::new(RwLock::new(HashMap::new())),
    57→        }
    58→    }
    59→
    60→    /// Register a new agent and create a session
    61→    pub async fn register_agent(&self, agent_id: AgentId, agent_name: String) -> (SessionId, FrameSender, InputReceiver) {
    62→        let session_id = Uuid::new_v4();
    63→
    64→        // Create channels
    65→        let (frame_tx, _) = broadcast::channel(16); // Buffer 16 frames
    66→        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64); // Buffer 64 input events
    67→
    68→        let session = Session {
    69→            id: session_id,
    70→            agent_id: agent_id.clone(),
    71→            agent_name,
    72→            started_at: chrono::Utc::now(),
    73→            viewer_count: 0,
    74→        };
    75→
    76→        let session_data = SessionData {
    77→            info: session,
    78→            frame_tx: frame_tx.clone(),
    79→            input_tx,
    80→            input_rx: None, // Will be taken by the agent handler
    81→        };
    82→
    83→        let mut sessions = self.sessions.write().await;
    84→        sessions.insert(session_id, session_data);
    85→
    86→        let mut agents = self.agents.write().await;
    87→        agents.insert(agent_id, session_id);
    88→
    89→        (session_id, frame_tx, input_rx)
    90→    }
    91→
    92→    /// Get a session by agent ID
    93→    pub async fn get_session_by_agent(&self, agent_id: &str) -> Option<Session> {
    94→        let agents = self.agents.read().await;
    95→        let session_id = agents.get(agent_id)?;
    96→
    97→        let sessions = self.sessions.read().await;
    98→        sessions.get(session_id).map(|s| s.info.clone())
    99→    }
   100→
   101→    /// Get a session by session ID
   102→    pub async fn get_session(&self, session_id: SessionId) -> Option<Session> {
   103→        let sessions = self.sessions.read().await;
   104→        sessions.get(&session_id).map(|s| s.info.clone())
   105→    }
   106→
   107→    /// Join a session as a viewer
   108→    pub async fn join_session(&self, session_id: SessionId) -> Option<(FrameReceiver, InputSender)> {
   109→        let mut sessions = self.sessions.write().await;
   110→        let session_data = sessions.get_mut(&session_id)?;
   111→
   112→        session_data.info.viewer_count += 1;
   113→
   114→        let frame_rx = session_data.frame_tx.subscribe();
   115→        let input_tx = session_data.input_tx.clone();
   116→
   117→        Some((frame_rx, input_tx))
   118→    }
   119→
   120→    /// Leave a session as a viewer
   121→    pub async fn leave_session(&self, session_id: SessionId) {
   122→        let mut sessions = self.sessions.write().await;
   123→        if let Some(session_data) = sessions.get_mut(&session_id) {
   124→            session_data.info.viewer_count = session_data.info.viewer_count.saturating_sub(1);
   125→        }
   126→    }
   127→
   128→    /// Remove a session (when agent disconnects)
   129→    pub async fn remove_session(&self, session_id: SessionId) {
   130→        let mut sessions = self.sessions.write().await;
   131→        if let Some(session_data) = sessions.remove(&session_id) {
   132→            let mut agents = self.agents.write().await;
   133→            agents.remove(&session_data.info.agent_id);
   134→        }
   135→    }
   136→
   137→    /// Disconnect a session by sending a disconnect message to the agent
   138→    /// Returns true if the message was sent successfully
   139→    pub async fn disconnect_session(&self, session_id: SessionId, reason: &str) -> bool {
   140→        let sessions = self.sessions.read().await;
   141→        if let Some(session_data) = sessions.get(&session_id) {
   142→            // Create disconnect message
   143→            use crate::proto;
   144→            use prost::Message;
   145→
   146→            let disconnect_msg = proto::Message {
   147→                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
   148→                    reason: reason.to_string(),
   149→                })),
   150→            };
   151→
   152→            let mut buf = Vec::new();
   153→            if disconnect_msg.encode(&mut buf).is_ok() {
   154→                // Send via input channel (will be forwarded to agent's WebSocket)
   155→                if session_data.input_tx.send(buf).await.is_ok() {
   156→                    return true;
   157→                }
   158→            }
   159→        }
   160→        false
   161→    }
   162→
   163→    /// List all active sessions
   164→    pub async fn list_sessions(&self) -> Vec<Session> {
   165→        let sessions = self.sessions.read().await;
   166→        sessions.values().map(|s| s.info.clone()).collect()
   167→    }
   168→}
   169→
   170→impl Default for SessionManager {
   171→    fn default() -> Self {
   172→        Self::new()
   173→    }
   174→}
   175→

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

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

use std::collections::HashMap;
use std::sync::Arc;
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;

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

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

/// 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 session = Session {
            id: session_id,
            agent_id: agent_id.clone(),
            agent_name,
            started_at: chrono::Utc::now(),
            viewer_count: 0,
        };

        let session_data = SessionData {
            info: session,
            frame_tx: frame_tx.clone(),
            input_tx,
            input_rx: None, // Will be taken by the agent handler
        };

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

    /// 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
    pub async fn join_session(&self, session_id: SessionId) -> 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 frame_rx = session_data.frame_tx.subscribe();
        let input_tx = session_data.input_tx.clone();

        Some((frame_rx, input_tx))
    }

    /// Leave a session as a viewer
    pub async fn leave_session(&self, session_id: SessionId) {
        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);
        }
    }

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

<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 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: {}", 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 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 frames from agent and broadcast to viewers
    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
                                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());
                            }
                            _ => {}
                        }
                    }
                    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;
        }
    };

    // Join the session
    let (mut frame_rx, input_tx) = match sessions.join_session(session_id).await {
        Some(channels) => channels,
        None => {
            warn!("Session not found: {}", session_id);
            return;
        }
    };

    info!("Viewer joined session: {}", 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();

    // 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(_)) => {
                                // 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: {}", session_id);
                break;
            }
            Ok(_) => {}
            Err(e) => {
                error!("WebSocket error from viewer: {}", e);
                break;
            }
        }
    }

    // Cleanup
    frame_forward.abort();
    sessions_cleanup.leave_session(session_id).await;
    info!("Viewer left session: {}", session_id);
}


Now I understand the server structure. I need to:

1. **Session manager changes:**
   - Store agent status (last heartbeat time, is_streaming)
   - Add methods to send StartStream/StopStream messages
   - Track which viewers are connected to which sessions

2. **Relay module changes:**
   - Handle AgentStatus and Heartbeat/HeartbeatAck from agent
   - When viewer connects, send StartStream to agent
   - When viewer disconnects (and no other viewers), send StopStream to agent
   - Handle heartbeat timeouts (detect zombie connections)

Let me update the session manager first with the new fields and methods.


TodoWrite

Add StartStream/StopStream/Heartbeat messages to protobuf

Implement agent state machine (idle vs streaming)

Add heartbeat mechanism with timeout detection

Update server to route viewer connections and send start/stop

Build and deploy updated server and agent


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 StartStream/StopStream/Heartbeat messages to protobuf

Implement agent state machine (idle vs streaming)

Add heartbeat mechanism with timeout detection

Update server to route viewer connections and send start/stop

Build and deploy updated server and agent


Now let me update the server's session manager to track streaming state and viewer IDs, and add methods for sending StartStream/StopStream.


Now let me update the server's session manager to track streaming state and add start/stop methods:


Write

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