     1→//! Agent configuration management
     2→//!
     3→//! Supports three configuration sources (in priority order):
     4→//! 1. Embedded config (magic bytes appended to executable)
     5→//! 2. Config file (guruconnect.toml or %ProgramData%\GuruConnect\agent.toml)
     6→//! 3. Environment variables (fallback)
     7→
     8→use anyhow::{anyhow, Context, Result};
     9→use serde::{Deserialize, Serialize};
    10→use std::io::{Read, Seek, SeekFrom};
    11→use std::path::PathBuf;
    12→use tracing::{info, warn};
    13→use uuid::Uuid;
    14→
    15→/// Magic marker for embedded configuration (10 bytes)
    16→const MAGIC_MARKER: &[u8] = b"GURUCONFIG";
    17→
    18→/// Embedded configuration data (appended to executable)
    19→#[derive(Debug, Clone, Serialize, Deserialize)]
    20→pub struct EmbeddedConfig {
    21→    /// Server WebSocket URL
    22→    pub server_url: String,
    23→    /// API key for authentication
    24→    pub api_key: String,
    25→    /// Company/organization name
    26→    #[serde(default)]
    27→    pub company: Option<String>,
    28→    /// Site/location name
    29→    #[serde(default)]
    30→    pub site: Option<String>,
    31→    /// Tags for categorization
    32→    #[serde(default)]
    33→    pub tags: Vec<String>,
    34→}
    35→
    36→/// Detected run mode based on filename
    37→#[derive(Debug, Clone, PartialEq)]
    38→pub enum RunMode {
    39→    /// Viewer-only installation (filename contains "Viewer")
    40→    Viewer,
    41→    /// Temporary support session (filename contains 6-digit code)
    42→    TempSupport(String),
    43→    /// Permanent agent with embedded config
    44→    PermanentAgent,
    45→    /// Unknown/default mode
    46→    Default,
    47→}
    48→
    49→/// Agent configuration
    50→#[derive(Debug, Clone, Serialize, Deserialize)]
    51→pub struct Config {
    52→    /// Server WebSocket URL (e.g., wss://connect.example.com/ws)
    53→    pub server_url: String,
    54→
    55→    /// Agent API key for authentication
    56→    pub api_key: String,
    57→
    58→    /// Unique agent identifier (generated on first run)
    59→    #[serde(default = "generate_agent_id")]
    60→    pub agent_id: String,
    61→
    62→    /// Optional hostname override
    63→    pub hostname_override: Option<String>,
    64→
    65→    /// Company/organization name (from embedded config)
    66→    #[serde(default)]
    67→    pub company: Option<String>,
    68→
    69→    /// Site/location name (from embedded config)
    70→    #[serde(default)]
    71→    pub site: Option<String>,
    72→
    73→    /// Tags for categorization (from embedded config)
    74→    #[serde(default)]
    75→    pub tags: Vec<String>,
    76→
    77→    /// Support code for one-time support sessions (set via command line or filename)
    78→    #[serde(skip)]
    79→    pub support_code: Option<String>,
    80→
    81→    /// Capture settings
    82→    #[serde(default)]
    83→    pub capture: CaptureConfig,
    84→
    85→    /// Encoding settings
    86→    #[serde(default)]
    87→    pub encoding: EncodingConfig,
    88→}
    89→
    90→fn generate_agent_id() -> String {
    91→    Uuid::new_v4().to_string()
    92→}
    93→
    94→#[derive(Debug, Clone, Serialize, Deserialize)]
    95→pub struct CaptureConfig {
    96→    /// Target frames per second (1-60)
    97→    #[serde(default = "default_fps")]
    98→    pub fps: u32,
    99→
   100→    /// Use DXGI Desktop Duplication (recommended)
   101→    #[serde(default = "default_true")]
   102→    pub use_dxgi: bool,
   103→
   104→    /// Fall back to GDI if DXGI fails
   105→    #[serde(default = "default_true")]
   106→    pub gdi_fallback: bool,
   107→}
   108→
   109→#[derive(Debug, Clone, Serialize, Deserialize)]
   110→pub struct EncodingConfig {
   111→    /// Preferred codec (auto, raw, vp9, h264)
   112→    #[serde(default = "default_codec")]
   113→    pub codec: String,
   114→
   115→    /// Quality (1-100, higher = better quality, more bandwidth)
   116→    #[serde(default = "default_quality")]
   117→    pub quality: u32,
   118→
   119→    /// Use hardware encoding if available
   120→    #[serde(default = "default_true")]
   121→    pub hardware_encoding: bool,
   122→}
   123→
   124→fn default_fps() -> u32 {
   125→    30
   126→}
   127→
   128→fn default_true() -> bool {
   129→    true
   130→}
   131→
   132→fn default_codec() -> String {
   133→    "auto".to_string()
   134→}
   135→
   136→fn default_quality() -> u32 {
   137→    75
   138→}
   139→
   140→impl Default for CaptureConfig {
   141→    fn default() -> Self {
   142→        Self {
   143→            fps: default_fps(),
   144→            use_dxgi: true,
   145→            gdi_fallback: true,
   146→        }
   147→    }
   148→}
   149→
   150→impl Default for EncodingConfig {
   151→    fn default() -> Self {
   152→        Self {
   153→            codec: default_codec(),
   154→            quality: default_quality(),
   155→            hardware_encoding: true,
   156→        }
   157→    }
   158→}
   159→
   160→impl Config {
   161→    /// Detect run mode from executable filename
   162→    pub fn detect_run_mode() -> RunMode {
   163→        let exe_path = match std::env::current_exe() {
   164→            Ok(p) => p,
   165→            Err(_) => return RunMode::Default,
   166→        };
   167→
   168→        let filename = match exe_path.file_stem() {
   169→            Some(s) => s.to_string_lossy().to_string(),
   170→            None => return RunMode::Default,
   171→        };
   172→
   173→        let filename_lower = filename.to_lowercase();
   174→
   175→        // Check for viewer mode
   176→        if filename_lower.contains("viewer") {
   177→            info!("Detected viewer mode from filename: {}", filename);
   178→            return RunMode::Viewer;
   179→        }
   180→
   181→        // Check for support code in filename (6-digit number)
   182→        if let Some(code) = Self::extract_support_code(&filename) {
   183→            info!("Detected support code from filename: {}", code);
   184→            return RunMode::TempSupport(code);
   185→        }
   186→
   187→        // Check for embedded config
   188→        if Self::has_embedded_config() {
   189→            info!("Detected embedded config in executable");
   190→            return RunMode::PermanentAgent;
   191→        }
   192→
   193→        RunMode::Default
   194→    }
   195→
   196→    /// Extract 6-digit support code from filename
   197→    fn extract_support_code(filename: &str) -> Option<String> {
   198→        // Look for patterns like "GuruConnect-123456" or "GuruConnect_123456"
   199→        for part in filename.split(|c| c == '-' || c == '_' || c == '.') {
   200→            let trimmed = part.trim();
   201→            if trimmed.len() == 6 && trimmed.chars().all(|c| c.is_ascii_digit()) {
   202→                return Some(trimmed.to_string());
   203→            }
   204→        }
   205→
   206→        // Check if last 6 characters are all digits
   207→        if filename.len() >= 6 {
   208→            let last_six = &filename[filename.len() - 6..];
   209→            if last_six.chars().all(|c| c.is_ascii_digit()) {
   210→                return Some(last_six.to_string());
   211→            }
   212→        }
   213→
   214→        None
   215→    }
   216→
   217→    /// Check if embedded configuration exists in the executable
   218→    pub fn has_embedded_config() -> bool {
   219→        Self::read_embedded_config().is_ok()
   220→    }
   221→
   222→    /// Read embedded configuration from the executable
   223→    pub fn read_embedded_config() -> Result<EmbeddedConfig> {
   224→        let exe_path = std::env::current_exe()
   225→            .context("Failed to get current executable path")?;
   226→
   227→        let mut file = std::fs::File::open(&exe_path)
   228→            .context("Failed to open executable for reading")?;
   229→
   230→        let file_size = file.metadata()?.len();
   231→        if file_size < (MAGIC_MARKER.len() + 4) as u64 {
   232→            return Err(anyhow!("File too small to contain embedded config"));
   233→        }
   234→
   235→        // Read the last part of the file to find magic marker
   236→        // Structure: [PE binary][GURUCONFIG][length:u32][json config]
   237→        // We need to search backwards from the end
   238→
   239→        // Read last 64KB (should be more than enough for config)
   240→        let search_size = std::cmp::min(65536, file_size as usize);
   241→        let search_start = file_size - search_size as u64;
   242→
   243→        file.seek(SeekFrom::Start(search_start))?;
   244→        let mut buffer = vec![0u8; search_size];
   245→        file.read_exact(&mut buffer)?;
   246→
   247→        // Find magic marker
   248→        let marker_pos = buffer.windows(MAGIC_MARKER.len())
   249→            .rposition(|window| window == MAGIC_MARKER)
   250→            .ok_or_else(|| anyhow!("Magic marker not found"))?;
   251→
   252→        // Read config length (4 bytes after marker)
   253→        let length_start = marker_pos + MAGIC_MARKER.len();
   254→        if length_start + 4 > buffer.len() {
   255→            return Err(anyhow!("Invalid embedded config: length field truncated"));
   256→        }
   257→
   258→        let config_length = u32::from_le_bytes([
   259→            buffer[length_start],
   260→            buffer[length_start + 1],
   261→            buffer[length_start + 2],
   262→            buffer[length_start + 3],
   263→        ]) as usize;
   264→
   265→        // Read config data
   266→        let config_start = length_start + 4;
   267→        if config_start + config_length > buffer.len() {
   268→            return Err(anyhow!("Invalid embedded config: data truncated"));
   269→        }
   270→
   271→        let config_bytes = &buffer[config_start..config_start + config_length];
   272→        let config: EmbeddedConfig = serde_json::from_slice(config_bytes)
   273→            .context("Failed to parse embedded config JSON")?;
   274→
   275→        info!("Loaded embedded config: server={}, company={:?}",
   276→              config.server_url, config.company);
   277→
   278→        Ok(config)
   279→    }
   280→
   281→    /// Check if an explicit agent configuration file exists
   282→    /// This returns true only if there's a real config file, not generated defaults
   283→    pub fn has_agent_config() -> bool {
   284→        // Check for embedded config first
   285→        if Self::has_embedded_config() {
   286→            return true;
   287→        }
   288→
   289→        // Check for config in current directory
   290→        let local_config = PathBuf::from("guruconnect.toml");
   291→        if local_config.exists() {
   292→            return true;
   293→        }
   294→
   295→        // Check in program data directory (Windows)
   296→        #[cfg(windows)]
   297→        {
   298→            if let Ok(program_data) = std::env::var("ProgramData") {
   299→                let path = PathBuf::from(program_data)
   300→                    .join("GuruConnect")
   301→                    .join("agent.toml");
   302→                if path.exists() {
   303→                    return true;
   304→                }
   305→            }
   306→        }
   307→
   308→        false
   309→    }
   310→
   311→    /// Load configuration from embedded config, file, or environment
   312→    pub fn load() -> Result<Self> {
   313→        // Priority 1: Try loading from embedded config
   314→        if let Ok(embedded) = Self::read_embedded_config() {
   315→            info!("Using embedded configuration");
   316→            let config = Config {
   317→                server_url: embedded.server_url,
   318→                api_key: embedded.api_key,
   319→                agent_id: generate_agent_id(),
   320→                hostname_override: None,
   321→                company: embedded.company,
   322→                site: embedded.site,
   323→                tags: embedded.tags,
   324→                support_code: None,
   325→                capture: CaptureConfig::default(),
   326→                encoding: EncodingConfig::default(),
   327→            };
   328→
   329→            // Save to file for persistence (so agent_id is preserved)
   330→            let _ = config.save();
   331→            return Ok(config);
   332→        }
   333→
   334→        // Priority 2: Try loading from config file
   335→        let config_path = Self::config_path();
   336→
   337→        if config_path.exists() {
   338→            let contents = std::fs::read_to_string(&config_path)
   339→                .with_context(|| format!("Failed to read config from {:?}", config_path))?;
   340→
   341→            let mut config: Config = toml::from_str(&contents)
   342→                .with_context(|| "Failed to parse config file")?;
   343→
   344→            // Ensure agent_id is set and saved
   345→            if config.agent_id.is_empty() {
   346→                config.agent_id = generate_agent_id();
   347→                let _ = config.save();
   348→            }
   349→
   350→            // support_code is always None when loading from file (set via CLI)
   351→            config.support_code = None;
   352→
   353→            return Ok(config);
   354→        }
   355→
   356→        // Priority 3: Fall back to environment variables
   357→        let server_url = std::env::var("GURUCONNECT_SERVER_URL")
   358→            .unwrap_or_else(|_| "wss://connect.azcomputerguru.com/ws/agent".to_string());
   359→
   360→        let api_key = std::env::var("GURUCONNECT_API_KEY")
   361→            .unwrap_or_else(|_| "dev-key".to_string());
   362→
   363→        let agent_id = std::env::var("GURUCONNECT_AGENT_ID")
   364→            .unwrap_or_else(|_| generate_agent_id());
   365→
   366→        let config = Config {
   367→            server_url,
   368→            api_key,
   369→            agent_id,
   370→            hostname_override: std::env::var("GURUCONNECT_HOSTNAME").ok(),
   371→            company: None,
   372→            site: None,
   373→            tags: Vec::new(),
   374→            support_code: None,
   375→            capture: CaptureConfig::default(),
   376→            encoding: EncodingConfig::default(),
   377→        };
   378→
   379→        // Save config with generated agent_id for persistence
   380→        let _ = config.save();
   381→
   382→        Ok(config)
   383→    }
   384→
   385→    /// Get the configuration file path
   386→    fn config_path() -> PathBuf {
   387→        // Check for config in current directory first
   388→        let local_config = PathBuf::from("guruconnect.toml");
   389→        if local_config.exists() {
   390→            return local_config;
   391→        }
   392→
   393→        // Check in program data directory (Windows)
   394→        #[cfg(windows)]
   395→        {
   396→            if let Ok(program_data) = std::env::var("ProgramData") {
   397→                let path = PathBuf::from(program_data)
   398→                    .join("GuruConnect")
   399→                    .join("agent.toml");
   400→                if path.exists() {
   401→                    return path;
   402→                }
   403→            }
   404→        }
   405→
   406→        // Default to local config
   407→        local_config
   408→    }
   409→
   410→    /// Get the hostname to use
   411→    pub fn hostname(&self) -> String {
   412→        self.hostname_override
   413→            .clone()
   414→            .unwrap_or_else(|| {
   415→                hostname::get()
   416→                    .map(|h| h.to_string_lossy().to_string())
   417→                    .unwrap_or_else(|_| "unknown".to_string())
   418→            })
   419→    }
   420→
   421→    /// Save current configuration to file
   422→    pub fn save(&self) -> Result<()> {
   423→        let config_path = Self::config_path();
   424→
   425→        // Ensure parent directory exists
   426→        if let Some(parent) = config_path.parent() {
   427→            std::fs::create_dir_all(parent)?;
   428→        }
   429→
   430→        let contents = toml::to_string_pretty(self)?;
   431→        std::fs::write(&config_path, contents)?;
   432→
   433→        Ok(())
   434→    }
   435→}
   436→
   437→/// Example configuration file content
   438→pub fn example_config() -> &'static str {
   439→    r#"# GuruConnect Agent Configuration
   440→
   441→# Server connection
   442→server_url = "wss://connect.example.com/ws"
   443→api_key = "your-agent-api-key"
   444→agent_id = "auto-generated-uuid"
   445→
   446→# Optional: override hostname
   447→# hostname_override = "custom-hostname"
   448→
   449→[capture]
   450→fps = 30
   451→use_dxgi = true
   452→gdi_fallback = true
   453→
   454→[encoding]
   455→codec = "auto"      # auto, raw, vp9, h264
   456→quality = 75        # 1-100
   457→hardware_encoding = true
   458→"#
   459→}
   460→

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

//! Agent configuration management
//!
//! Supports three configuration sources (in priority order):
//! 1. Embedded config (magic bytes appended to executable)
//! 2. Config file (guruconnect.toml or %ProgramData%\GuruConnect\agent.toml)
//! 3. Environment variables (fallback)

use anyhow::{anyhow, Context, Result};
use serde::{Deserialize, Serialize};
use std::io::{Read, Seek, SeekFrom};
use std::path::PathBuf;
use tracing::{info, warn};
use uuid::Uuid;

/// Magic marker for embedded configuration (10 bytes)
const MAGIC_MARKER: &[u8] = b"GURUCONFIG";

/// Embedded configuration data (appended to executable)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddedConfig {
    /// Server WebSocket URL
    pub server_url: String,
    /// API key for authentication
    pub api_key: String,
    /// Company/organization name
    #[serde(default)]
    pub company: Option<String>,
    /// Site/location name
    #[serde(default)]
    pub site: Option<String>,
    /// Tags for categorization
    #[serde(default)]
    pub tags: Vec<String>,
}

/// Detected run mode based on filename
#[derive(Debug, Clone, PartialEq)]
pub enum RunMode {
    /// Viewer-only installation (filename contains "Viewer")
    Viewer,
    /// Temporary support session (filename contains 6-digit code)
    TempSupport(String),
    /// Permanent agent with embedded config
    PermanentAgent,
    /// Unknown/default mode
    Default,
}

/// Agent configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Config {
    /// Server WebSocket URL (e.g., wss://connect.example.com/ws)
    pub server_url: String,

    /// Agent API key for authentication
    pub api_key: String,

    /// Unique agent identifier (generated on first run)
    #[serde(default = "generate_agent_id")]
    pub agent_id: String,

    /// Optional hostname override
    pub hostname_override: Option<String>,

    /// Company/organization name (from embedded config)
    #[serde(default)]
    pub company: Option<String>,

    /// Site/location name (from embedded config)
    #[serde(default)]
    pub site: Option<String>,

    /// Tags for categorization (from embedded config)
    #[serde(default)]
    pub tags: Vec<String>,

    /// Support code for one-time support sessions (set via command line or filename)
    #[serde(skip)]
    pub support_code: Option<String>,

    /// Capture settings
    #[serde(default)]
    pub capture: CaptureConfig,

    /// Encoding settings
    #[serde(default)]
    pub encoding: EncodingConfig,
}

fn generate_agent_id() -> String {
    Uuid::new_v4().to_string()
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CaptureConfig {
    /// Target frames per second (1-60)
    #[serde(default = "default_fps")]
    pub fps: u32,

    /// Use DXGI Desktop Duplication (recommended)
    #[serde(default = "default_true")]
    pub use_dxgi: bool,

    /// Fall back to GDI if DXGI fails
    #[serde(default = "default_true")]
    pub gdi_fallback: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncodingConfig {
    /// Preferred codec (auto, raw, vp9, h264)
    #[serde(default = "default_codec")]
    pub codec: String,

    /// Quality (1-100, higher = better quality, more bandwidth)
    #[serde(default = "default_quality")]
    pub quality: u32,

    /// Use hardware encoding if available
    #[serde(default = "default_true")]
    pub hardware_encoding: bool,
}

fn default_fps() -> u32 {
    30
}

fn default_true() -> bool {
    true
}

fn default_codec() -> String {
    "auto".to_string()
}

fn default_quality() -> u32 {
    75
}

impl Default for CaptureConfig {
    fn default() -> Self {
        Self {
            fps: default_fps(),
            use_dxgi: true,
            gdi_fallback: true,
        }
    }
}

impl Default for EncodingConfig {
    fn default() -> Self {
        Self {
            codec: default_codec(),
            quality: default_quality(),
            hardware_encoding: true,
        }
    }
}

impl Config {
    /// Detect run mode from executable filename
    pub fn detect_run_mode() -> RunMode {
        let exe_path = match std::env::current_exe() {
            Ok(p) => p,
            Err(_) => return RunMode::Default,
        };

        let filename = match exe_path.file_stem() {
            Some(s) => s.to_string_lossy().to_string(),
            None => return RunMode::Default,
        };

        let filename_lower = filename.to_lowercase();

        // Check for viewer mode
        if filename_lower.contains("viewer") {
            info!("Detected viewer mode from filename: {}", filename);
            return RunMode::Viewer;
        }

        // Check for support code in filename (6-digit number)
        if let Some(code) = Self::extract_support_code(&filename) {
            info!("Detected support code from filename: {}", code);
            return RunMode::TempSupport(code);
        }

        // Check for embedded config
        if Self::has_embedded_config() {
            info!("Detected embedded config in executable");
            return RunMode::PermanentAgent;
        }

        RunMode::Default
    }

    /// Extract 6-digit support code from filename
    fn extract_support_code(filename: &str) -> Option<String> {
        // Look for patterns like "GuruConnect-123456" or "GuruConnect_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());
            }
        }

        // Check if last 6 characters are all digits
        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 embedded configuration exists in the executable
    pub fn has_embedded_config() -> bool {
        Self::read_embedded_config().is_ok()
    }

    /// Read embedded configuration from the executable
    pub fn read_embedded_config() -> Result<EmbeddedConfig> {
        let exe_path = std::env::current_exe()
            .context("Failed to get current executable path")?;

        let mut file = std::fs::File::open(&exe_path)
            .context("Failed to open executable for reading")?;

        let file_size = file.metadata()?.len();
        if file_size < (MAGIC_MARKER.len() + 4) as u64 {
            return Err(anyhow!("File too small to contain embedded config"));
        }

        // Read the last part of the file to find magic marker
        // Structure: [PE binary][GURUCONFIG][length:u32][json config]
        // We need to search backwards from the end

        // Read last 64KB (should be more than enough for config)
        let search_size = std::cmp::min(65536, file_size as usize);
        let search_start = file_size - search_size as u64;

        file.seek(SeekFrom::Start(search_start))?;
        let mut buffer = vec![0u8; search_size];
        file.read_exact(&mut buffer)?;

        // Find magic marker
        let marker_pos = buffer.windows(MAGIC_MARKER.len())
            .rposition(|window| window == MAGIC_MARKER)
            .ok_or_else(|| anyhow!("Magic marker not found"))?;

        // Read config length (4 bytes after marker)
        let length_start = marker_pos + MAGIC_MARKER.len();
        if length_start + 4 > buffer.len() {
            return Err(anyhow!("Invalid embedded config: length field truncated"));
        }

        let config_length = u32::from_le_bytes([
            buffer[length_start],
            buffer[length_start + 1],
            buffer[length_start + 2],
            buffer[length_start + 3],
        ]) as usize;

        // Read config data
        let config_start = length_start + 4;
        if config_start + config_length > buffer.len() {
            return Err(anyhow!("Invalid embedded config: data truncated"));
        }

        let config_bytes = &buffer[config_start..config_start + config_length];
        let config: EmbeddedConfig = serde_json::from_slice(config_bytes)
            .context("Failed to parse embedded config JSON")?;

        info!("Loaded embedded config: server={}, company={:?}",
              config.server_url, config.company);

        Ok(config)
    }

    /// Check if an explicit agent configuration file exists
    /// This returns true only if there's a real config file, not generated defaults
    pub fn has_agent_config() -> bool {
        // Check for embedded config first
        if Self::has_embedded_config() {
            return true;
        }

        // Check for config in current directory
        let local_config = PathBuf::from("guruconnect.toml");
        if local_config.exists() {
            return true;
        }

        // Check in program data directory (Windows)
        #[cfg(windows)]
        {
            if let Ok(program_data) = std::env::var("ProgramData") {
                let path = PathBuf::from(program_data)
                    .join("GuruConnect")
                    .join("agent.toml");
                if path.exists() {
                    return true;
                }
            }
        }

        false
    }

    /// Load configuration from embedded config, file, or environment
    pub fn load() -> Result<Self> {
        // Priority 1: Try loading from embedded config
        if let Ok(embedded) = Self::read_embedded_config() {
            info!("Using embedded configuration");
            let config = Config {
                server_url: embedded.server_url,
                api_key: embedded.api_key,
                agent_id: generate_agent_id(),
                hostname_override: None,
                company: embedded.company,
                site: embedded.site,
                tags: embedded.tags,
                support_code: None,
                capture: CaptureConfig::default(),
                encoding: EncodingConfig::default(),
            };

            // Save to file for persistence (so agent_id is preserved)
            let _ = config.save();
            return Ok(config);
        }

        // Priority 2: Try loading from config file
        let config_path = Self::config_path();

        if config_path.exists() {
            let contents = std::fs::read_to_string(&config_path)
                .with_context(|| format!("Failed to read config from {:?}", config_path))?;

            let mut config: Config = toml::from_str(&contents)
                .with_context(|| "Failed to parse config file")?;

            // Ensure agent_id is set and saved
            if config.agent_id.is_empty() {
                config.agent_id = generate_agent_id();
                let _ = config.save();
            }

            // support_code is always None when loading from file (set via CLI)
            config.support_code = None;

            return Ok(config);
        }

        // Priority 3: Fall back to environment variables
        let server_url = std::env::var("GURUCONNECT_SERVER_URL")
            .unwrap_or_else(|_| "wss://connect.azcomputerguru.com/ws/agent".to_string());

        let api_key = std::env::var("GURUCONNECT_API_KEY")
            .unwrap_or_else(|_| "dev-key".to_string());

        let agent_id = std::env::var("GURUCONNECT_AGENT_ID")
            .unwrap_or_else(|_| generate_agent_id());

        let config = Config {
            server_url,
            api_key,
            agent_id,
            hostname_override: std::env::var("GURUCONNECT_HOSTNAME").ok(),
            company: None,
            site: None,
            tags: Vec::new(),
            support_code: None,
            capture: CaptureConfig::default(),
            encoding: EncodingConfig::default(),
        };

        // Save config with generated agent_id for persistence
        let _ = config.save();

        Ok(config)
    }

    /// Get the configuration file path
    fn config_path() -> PathBuf {
        // Check for config in current directory first
        let local_config = PathBuf::from("guruconnect.toml");
        if local_config.exists() {
            return local_config;
        }

        // Check in program data directory (Windows)
        #[cfg(windows)]
        {
            if let Ok(program_data) = std::env::var("ProgramData") {
                let path = PathBuf::from(program_data)
                    .join("GuruConnect")
                    .join("agent.toml");
                if path.exists() {
                    return path;
                }
            }
        }

        // Default to local config
        local_config
    }

    /// Get the hostname to use
    pub fn hostname(&self) -> String {
        self.hostname_override
            .clone()
            .unwrap_or_else(|| {
                hostname::get()
                    .map(|h| h.to_string_lossy().to_string())
                    .unwrap_or_else(|_| "unknown".to_string())
            })
    }

    /// Save current configuration to file
    pub fn save(&self) -> Result<()> {
        let config_path = Self::config_path();

        // Ensure parent directory exists
        if let Some(parent) = config_path.parent() {
            std::fs::create_dir_all(parent)?;
        }

        let contents = toml::to_string_pretty(self)?;
        std::fs::write(&config_path, contents)?;

        Ok(())
    }
}

/// Example configuration file content
pub fn example_config() -> &'static str {
    r#"# GuruConnect Agent Configuration

# Server connection
server_url = "wss://connect.example.com/ws"
api_key = "your-agent-api-key"
agent_id = "auto-generated-uuid"

# Optional: override hostname
# hostname_override = "custom-hostname"

[capture]
fps = 30
use_dxgi = true
gdi_fallback = true

[encoding]
codec = "auto"      # auto, raw, vp9, h264
quality = 75        # 1-100
hardware_encoding = true
"#
}


I see the config supports `use_dxgi` and `gdi_fallback` settings. The issue is that embedded config doesn't carry those capture settings - it uses defaults. Let me add panic handling and more verbose logging to init_streaming, and also force GDI mode for this Geovision test.

Let me update the session module to:
1. Add a panic handler using std::panic::catch_unwind
2. Log more details about what's happening during init_streaming
3. Force GDI mode on this agent for testing

Actually, I should create a new custom agent build that forces GDI mode. The embedded config currently uses `CaptureConfig::default()` which has `use_dxgi: true`.

Let me modify the embedded config parsing to support capture settings, or better yet, let me modify the EmbeddedConfig struct to include capture settings that can override defaults.

Wait, the simpler approach is to just modify the init_streaming to try GDI first for debugging, or add panic catching. Let me add better error handling first.

Actually, the simplest fix is to:
1. Add catch_unwind around init_streaming components
2. Make sure DXGI errors are properly caught and fallback works
3. Add more logging

Let me do this properly.


The embedded config doesn't include capture settings, so it defaults to DXGI. The DXGI initialization might be panicking on the Geovision system. Let me add panic handling and force GDI mode for testing:


TodoWrite

Add panic handling to init_streaming to prevent agent crash

Create custom agent with GDI-only mode for Geovision testing

Deploy and test on DESKTOP-I66IM5Q


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 panic handling to init_streaming to prevent agent crash

Create custom agent with GDI-only mode for Geovision testing

Deploy and test on DESKTOP-I66IM5Q


Now let me modify the agent's init_streaming to use catch_unwind to prevent panics from crashing the entire process, and add more logging. I'll also add an option to force GDI mode.


Edit

    /// Initialize streaming resources (capturer, encoder, input)
    fn init_streaming(&mut self) -> Result<()> {
        if self.capturer.is_some() {
            return Ok(()); // Already initialized
        }

        tracing::info!("Initializing streaming resources...");

        // Get primary display
        let primary_display = capture::primary_display()?;
        tracing::info!("Using display: {} ({}x{})",
            primary_display.name, primary_display.width, primary_display.height);

        // Create capturer
        let capturer = capture::create_capturer(
            primary_display.clone(),
            self.config.capture.use_dxgi,
            self.config.capture.gdi_fallback,
        )?;
        self.capturer = Some(capturer);

        // Create encoder
        let encoder = encoder::create_encoder(
            &self.config.encoding.codec,
            self.config.encoding.quality,
        )?;
        self.encoder = Some(encoder);

        // Create input controller
        let input = InputController::new()?;
        self.input = Some(input);

        tracing::info!("Streaming resources initialized");
        Ok(())
    }

    /// Initialize streaming resources (capturer, encoder, input)
    fn init_streaming(&mut self) -> Result<()> {
        if self.capturer.is_some() {
            return Ok(()); // Already initialized
        }

        tracing::info!("Initializing streaming resources...");
        tracing::info!("Capture config: use_dxgi={}, gdi_fallback={}, fps={}",
            self.config.capture.use_dxgi, self.config.capture.gdi_fallback, self.config.capture.fps);

        // Get primary display with panic protection
        tracing::debug!("Enumerating displays...");
        let primary_display = match std::panic::catch_unwind(|| capture::primary_display()) {
            Ok(result) => result?,
            Err(e) => {
                tracing::error!("Panic during display enumeration: {:?}", e);
                return Err(anyhow::anyhow!("Display enumeration panicked"));
            }
        };
        tracing::info!("Using display: {} ({}x{})",
            primary_display.name, primary_display.width, primary_display.height);

        // Create capturer with panic protection
        // Force GDI mode if DXGI fails or panics
        tracing::debug!("Creating capturer (DXGI={})...", self.config.capture.use_dxgi);
        let capturer = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            capture::create_capturer(
                primary_display.clone(),
                self.config.capture.use_dxgi,
                self.config.capture.gdi_fallback,
            )
        })) {
            Ok(result) => result?,
            Err(e) => {
                tracing::error!("Panic during capturer creation: {:?}", e);
                // Try GDI-only as last resort
                tracing::warn!("Attempting GDI-only capture after DXGI panic...");
                capture::create_capturer(primary_display.clone(), false, false)?
            }
        };
        self.capturer = Some(capturer);
        tracing::info!("Capturer created successfully");

        // Create encoder with panic protection
        tracing::debug!("Creating encoder (codec={}, quality={})...",
            self.config.encoding.codec, self.config.encoding.quality);
        let encoder = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            encoder::create_encoder(
                &self.config.encoding.codec,
                self.config.encoding.quality,
            )
        })) {
            Ok(result) => result?,
            Err(e) => {
                tracing::error!("Panic during encoder creation: {:?}", e);
                return Err(anyhow::anyhow!("Encoder creation panicked"));
            }
        };
        self.encoder = Some(encoder);
        tracing::info!("Encoder created successfully");

        // Create input controller with panic protection
        tracing::debug!("Creating input controller...");
        let input = match std::panic::catch_unwind(InputController::new) {
            Ok(result) => result?,
            Err(e) => {
                tracing::error!("Panic during input controller creation: {:?}", e);
                return Err(anyhow::anyhow!("Input controller creation panicked"));
            }
        };
        self.input = Some(input);

        tracing::info!("Streaming resources initialized successfully");
        Ok(())
    }


Error: File has been unexpectedly modified. Read it again before attempting to write it.


The file was modified, need to re-read it first.


   116→    /// Initialize streaming resources (capturer, encoder, input)
   117→    fn init_streaming(&mut self) -> Result<()> {
   118→        if self.capturer.is_some() {
   119→            return Ok(()); // Already initialized
   120→        }
   121→
   122→        tracing::info!("Initializing streaming resources...");
   123→
   124→        // Get primary display
   125→        let primary_display = capture::primary_display()?;
   126→        tracing::info!("Using display: {} ({}x{})",
   127→            primary_display.name, primary_display.width, primary_display.height);
   128→
   129→        // Create capturer
   130→        let capturer = capture::create_capturer(
   131→            primary_display.clone(),
   132→            self.config.capture.use_dxgi,
   133→            self.config.capture.gdi_fallback,
   134→        )?;
   135→        self.capturer = Some(capturer);
   136→
   137→        // Create encoder
   138→        let encoder = encoder::create_encoder(
   139→            &self.config.encoding.codec,
   140→            self.config.encoding.quality,
   141→        )?;
   142→        self.encoder = Some(encoder);
   143→
   144→        // Create input controller
   145→        let input = InputController::new()?;
   146→        self.input = Some(input);
   147→
   148→        tracing::info!("Streaming resources initialized");
   149→        Ok(())
   150→    }
   151→
   152→    /// Release streaming resources to save CPU/memory when idle
   153→    fn release_streaming(&mut self) {
   154→        if self.capturer.is_some() {
   155→            tracing::info!("Releasing streaming resources");
   156→            self.capturer = None;
   157→            self.encoder = None;
   158→            self.input = None;
   159→            self.current_viewer_id = None;
   160→        }
   161→    }
   162→
   163→    /// Get display count for status reports
   164→    fn get_display_count(&self) -> i32 {
   165→        capture::enumerate_displays().map(|d| d.len() as i32).unwrap_or(1)

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

    /// Initialize streaming resources (capturer, encoder, input)
    fn init_streaming(&mut self) -> Result<()> {
        if self.capturer.is_some() {
            return Ok(()); // Already initialized
        }

        tracing::info!("Initializing streaming resources...");

        // Get primary display
        let primary_display = capture::primary_display()?;
        tracing::info!("Using display: {} ({}x{})",
            primary_display.name, primary_display.width, primary_display.height);

        // Create capturer
        let capturer = capture::create_capturer(
            primary_display.clone(),
            self.config.capture.use_dxgi,
            self.config.capture.gdi_fallback,
        )?;
        self.capturer = Some(capturer);

        // Create encoder
        let encoder = encoder::create_encoder(
            &self.config.encoding.codec,
            self.config.encoding.quality,
        )?;
        self.encoder = Some(encoder);

        // Create input controller
        let input = InputController::new()?;
        self.input = Some(input);

        tracing::info!("Streaming resources initialized");
        Ok(())
    }

    /// Release streaming resources to save CPU/memory when idle
    fn release_streaming(&mut self) {
        if self.capturer.is_some() {
            tracing::info!("Releasing streaming resources");
            self.capturer = None;
            self.encoder = None;
            self.input = None;
            self.current_viewer_id = None;
        }
    }

    /// Get display count for status reports
    fn get_display_count(&self) -> i32 {
        capture::enumerate_displays().map(|d| d.len() as i32).unwrap_or(1)


Good, the file looks the same. Let me try the edit again:


    /// Initialize streaming resources (capturer, encoder, input)
    fn init_streaming(&mut self) -> Result<()> {
        if self.capturer.is_some() {
            return Ok(()); // Already initialized
        }

        tracing::info!("Initializing streaming resources...");

        // Get primary display
        let primary_display = capture::primary_display()?;
        tracing::info!("Using display: {} ({}x{})",
            primary_display.name, primary_display.width, primary_display.height);

        // Create capturer
        let capturer = capture::create_capturer(
            primary_display.clone(),
            self.config.capture.use_dxgi,
            self.config.capture.gdi_fallback,
        )?;
        self.capturer = Some(capturer);

        // Create encoder
        let encoder = encoder::create_encoder(
            &self.config.encoding.codec,
            self.config.encoding.quality,
        )?;
        self.encoder = Some(encoder);

        // Create input controller
        let input = InputController::new()?;
        self.input = Some(input);

        tracing::info!("Streaming resources initialized");
        Ok(())
    }

    /// Initialize streaming resources (capturer, encoder, input)
    fn init_streaming(&mut self) -> Result<()> {
        if self.capturer.is_some() {
            return Ok(()); // Already initialized
        }

        tracing::info!("Initializing streaming resources...");
        tracing::info!("Capture config: use_dxgi={}, gdi_fallback={}, fps={}",
            self.config.capture.use_dxgi, self.config.capture.gdi_fallback, self.config.capture.fps);

        // Get primary display with panic protection
        tracing::debug!("Enumerating displays...");
        let primary_display = match std::panic::catch_unwind(|| capture::primary_display()) {
            Ok(result) => result?,
            Err(e) => {
                tracing::error!("Panic during display enumeration: {:?}", e);
                return Err(anyhow::anyhow!("Display enumeration panicked"));
            }
        };
        tracing::info!("Using display: {} ({}x{})",
            primary_display.name, primary_display.width, primary_display.height);

        // Create capturer with panic protection
        // Force GDI mode if DXGI fails or panics
        tracing::debug!("Creating capturer (DXGI={})...", self.config.capture.use_dxgi);
        let capturer = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            capture::create_capturer(
                primary_display.clone(),
                self.config.capture.use_dxgi,
                self.config.capture.gdi_fallback,
            )
        })) {
            Ok(result) => result?,
            Err(e) => {
                tracing::error!("Panic during capturer creation: {:?}", e);
                // Try GDI-only as last resort
                tracing::warn!("Attempting GDI-only capture after DXGI panic...");
                capture::create_capturer(primary_display.clone(), false, false)?
            }
        };
        self.capturer = Some(capturer);
        tracing::info!("Capturer created successfully");

        // Create encoder with panic protection
        tracing::debug!("Creating encoder (codec={}, quality={})...",
            self.config.encoding.codec, self.config.encoding.quality);
        let encoder = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            encoder::create_encoder(
                &self.config.encoding.codec,
                self.config.encoding.quality,
            )
        })) {
            Ok(result) => result?,
            Err(e) => {
                tracing::error!("Panic during encoder creation: {:?}", e);
                return Err(anyhow::anyhow!("Encoder creation panicked"));
            }
        };
        self.encoder = Some(encoder);
        tracing::info!("Encoder created successfully");

        // Create input controller with panic protection
        tracing::debug!("Creating input controller...");
        let input = match std::panic::catch_unwind(InputController::new) {
            Ok(result) => result?,
            Err(e) => {
                tracing::error!("Panic during input controller creation: {:?}", e);
                return Err(anyhow::anyhow!("Input controller creation panicked"));
            }
        };
        self.input = Some(input);

        tracing::info!("Streaming resources initialized successfully");
        Ok(())
    }


The file keeps getting modified. Let me read the whole file again.


Bash

cd C:/Users/MikeSwanson/Claude/guru-connect && git diff agent/src/session/mod.rs | head -100

Check if file has uncommitted changes


No changes in git. Let me read the file again and use a longer timeout:


     1→//! Session management for the agent
     2→//!
     3→//! Handles the lifecycle of a remote session including:
     4→//! - Connection to server
     5→//! - Idle mode (heartbeat only, minimal resources)
     6→//! - Active/streaming mode (capture and send frames)
     7→//! - Input event handling
     8→
     9→#[cfg(windows)]
    10→use windows::Win32::System::Console::{AllocConsole, GetConsoleWindow};
    11→#[cfg(windows)]
    12→use windows::Win32::UI::WindowsAndMessaging::{ShowWindow, SW_SHOW};
    13→
    14→use crate::capture::{self, Capturer, Display};
    15→use crate::chat::{ChatController, ChatMessage as ChatMsg};
    16→use crate::config::Config;
    17→use crate::encoder::{self, Encoder};
    18→use crate::input::InputController;
    19→
    20→/// Show the debug console window (Windows only)
    21→#[cfg(windows)]
    22→fn show_debug_console() {
    23→    unsafe {
    24→        let hwnd = GetConsoleWindow();
    25→        if hwnd.0 == std::ptr::null_mut() {
    26→            let _ = AllocConsole();
    27→            tracing::info!("Debug console window opened");
    28→        } else {
    29→            let _ = ShowWindow(hwnd, SW_SHOW);
    30→            tracing::info!("Debug console window shown");
    31→        }
    32→    }
    33→}
    34→
    35→#[cfg(not(windows))]
    36→fn show_debug_console() {
    37→    // No-op on non-Windows platforms
    38→}
    39→
    40→use crate::proto::{Message, message, ChatMessage, AgentStatus, Heartbeat, HeartbeatAck};
    41→use crate::transport::WebSocketTransport;
    42→use crate::tray::{TrayController, TrayAction};
    43→use anyhow::Result;
    44→use std::time::{Duration, Instant};
    45→
    46→// Heartbeat interval (30 seconds)
    47→const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(30);
    48→// Status report interval (60 seconds)
    49→const STATUS_INTERVAL: Duration = Duration::from_secs(60);
    50→// Update check interval (1 hour)
    51→const UPDATE_CHECK_INTERVAL: Duration = Duration::from_secs(3600);
    52→
    53→/// Session manager handles the remote control session
    54→pub struct SessionManager {
    55→    config: Config,
    56→    transport: Option<WebSocketTransport>,
    57→    state: SessionState,
    58→    // Lazy-initialized streaming resources
    59→    capturer: Option<Box<dyn Capturer>>,
    60→    encoder: Option<Box<dyn Encoder>>,
    61→    input: Option<InputController>,
    62→    // Streaming state
    63→    current_viewer_id: Option<String>,
    64→    // System info for status reports
    65→    hostname: String,
    66→    is_elevated: bool,
    67→    start_time: Instant,
    68→}
    69→
    70→#[derive(Debug, Clone, PartialEq)]
    71→enum SessionState {
    72→    Disconnected,
    73→    Connecting,
    74→    Idle,       // Connected but not streaming - minimal resource usage
    75→    Streaming,  // Actively capturing and sending frames
    76→}
    77→
    78→impl SessionManager {
    79→    /// Create a new session manager
    80→    pub fn new(config: Config, is_elevated: bool) -> Self {
    81→        let hostname = config.hostname();
    82→        Self {
    83→            config,
    84→            transport: None,
    85→            state: SessionState::Disconnected,
    86→            capturer: None,
    87→            encoder: None,
    88→            input: None,
    89→            current_viewer_id: None,
    90→            hostname,
    91→            is_elevated,
    92→            start_time: Instant::now(),
    93→        }
    94→    }
    95→
    96→    /// Connect to the server
    97→    pub async fn connect(&mut self) -> Result<()> {
    98→        self.state = SessionState::Connecting;
    99→
   100→        let transport = WebSocketTransport::connect(
   101→            &self.config.server_url,
   102→            &self.config.agent_id,
   103→            &self.config.api_key,
   104→            Some(&self.hostname),
   105→            self.config.support_code.as_deref(),
   106→        ).await?;
   107→
   108→        self.transport = Some(transport);
   109→        self.state = SessionState::Idle;  // Start in idle mode
   110→
   111→        tracing::info!("Connected to server, entering idle mode");
   112→
   113→        Ok(())
   114→    }
   115→
   116→    /// Initialize streaming resources (capturer, encoder, input)
   117→    fn init_streaming(&mut self) -> Result<()> {
   118→        if self.capturer.is_some() {
   119→            return Ok(()); // Already initialized
   120→        }
   121→
   122→        tracing::info!("Initializing streaming resources...");
   123→
   124→        // Get primary display
   125→        let primary_display = capture::primary_display()?;
   126→        tracing::info!("Using display: {} ({}x{})",
   127→            primary_display.name, primary_display.width, primary_display.height);
   128→
   129→        // Create capturer
   130→        let capturer = capture::create_capturer(
   131→            primary_display.clone(),
   132→            self.config.capture.use_dxgi,
   133→            self.config.capture.gdi_fallback,
   134→        )?;
   135→        self.capturer = Some(capturer);
   136→
   137→        // Create encoder
   138→        let encoder = encoder::create_encoder(
   139→            &self.config.encoding.codec,
   140→            self.config.encoding.quality,
   141→        )?;
   142→        self.encoder = Some(encoder);
   143→
   144→        // Create input controller
   145→        let input = InputController::new()?;
   146→        self.input = Some(input);
   147→
   148→        tracing::info!("Streaming resources initialized");
   149→        Ok(())
   150→    }
   151→
   152→    /// Release streaming resources to save CPU/memory when idle
   153→    fn release_streaming(&mut self) {
   154→        if self.capturer.is_some() {
   155→            tracing::info!("Releasing streaming resources");
   156→            self.capturer = None;
   157→            self.encoder = None;
   158→            self.input = None;
   159→            self.current_viewer_id = None;
   160→        }
   161→    }
   162→
   163→    /// Get display count for status reports
   164→    fn get_display_count(&self) -> i32 {
   165→        capture::enumerate_displays().map(|d| d.len() as i32).unwrap_or(1)
   166→    }
   167→
   168→    /// Send agent status to server
   169→    async fn send_status(&mut self) -> Result<()> {
   170→        let status = AgentStatus {
   171→            hostname: self.hostname.clone(),
   172→            os_version: std::env::consts::OS.to_string(),
   173→            is_elevated: self.is_elevated,
   174→            uptime_secs: self.start_time.elapsed().as_secs() as i64,
   175→            display_count: self.get_display_count(),
   176→            is_streaming: self.state == SessionState::Streaming,
   177→            agent_version: crate::build_info::short_version(),
   178→            organization: self.config.company.clone().unwrap_or_default(),
   179→            site: self.config.site.clone().unwrap_or_default(),
   180→            tags: self.config.tags.clone(),
   181→        };
   182→
   183→        let msg = Message {
   184→            payload: Some(message::Payload::AgentStatus(status)),
   185→        };
   186→
   187→        if let Some(transport) = self.transport.as_mut() {
   188→            transport.send(msg).await?;
   189→        }
   190→
   191→        Ok(())
   192→    }
   193→
   194→    /// Send heartbeat to server
   195→    async fn send_heartbeat(&mut self) -> Result<()> {
   196→        let heartbeat = Heartbeat {
   197→            timestamp: chrono::Utc::now().timestamp_millis(),
   198→        };
   199→
   200→        let msg = Message {
   201→            payload: Some(message::Payload::Heartbeat(heartbeat)),
   202→        };
   203→
   204→        if let Some(transport) = self.transport.as_mut() {
   205→            transport.send(msg).await?;
   206→        }
   207→
   208→        Ok(())
   209→    }
   210→
   211→    /// Run the session main loop with tray and chat event processing
   212→    pub async fn run_with_tray(&mut self, tray: Option<&TrayController>, chat: Option<&ChatController>) -> Result<()> {
   213→        if self.transport.is_none() {
   214→            anyhow::bail!("Not connected");
   215→        }
   216→
   217→        // Send initial status
   218→        self.send_status().await?;
   219→
   220→        // Timing for heartbeat and status
   221→        let mut last_heartbeat = Instant::now();
   222→        let mut last_status = Instant::now();
   223→        let mut last_frame_time = Instant::now();
   224→        let mut last_update_check = Instant::now();
   225→        let frame_interval = Duration::from_millis(1000 / self.config.capture.fps as u64);
   226→
   227→        // Main loop
   228→        loop {
   229→            // Process tray events
   230→            if let Some(t) = tray {
   231→                if let Some(action) = t.process_events() {
   232→                    match action {
   233→                        TrayAction::EndSession => {
   234→                            tracing::info!("User requested session end via tray");
   235→                            return Err(anyhow::anyhow!("USER_EXIT: Session ended by user"));
   236→                        }
   237→                        TrayAction::ShowDetails => {
   238→                            tracing::info!("User requested details (not yet implemented)");
   239→                        }
   240→                        TrayAction::ShowDebugWindow => {
   241→                            show_debug_console();
   242→                        }
   243→                    }
   244→                }
   245→
   246→                if t.exit_requested() {
   247→                    tracing::info!("Exit requested via tray");
   248→                    return Err(anyhow::anyhow!("USER_EXIT: Exit requested by user"));
   249→                }
   250→            }
   251→
   252→            // Process incoming messages
   253→            let messages: Vec<Message> = {
   254→                let transport = self.transport.as_mut().unwrap();
   255→                let mut msgs = Vec::new();
   256→                while let Some(msg) = transport.try_recv()? {
   257→                    msgs.push(msg);
   258→                }
   259→                msgs
   260→            };
   261→
   262→            for msg in messages {
   263→                // Handle chat messages specially
   264→                if let Some(message::Payload::ChatMessage(chat_msg)) = &msg.payload {
   265→                    if let Some(c) = chat {
   266→                        c.add_message(ChatMsg {
   267→                            id: chat_msg.id.clone(),
   268→                            sender: chat_msg.sender.clone(),
   269→                            content: chat_msg.content.clone(),
   270→                            timestamp: chat_msg.timestamp,
   271→                        });
   272→                    }
   273→                    continue;
   274→                }
   275→
   276→                // Handle control messages that affect state
   277→                if let Some(ref payload) = msg.payload {
   278→                    match payload {
   279→                        message::Payload::StartStream(start) => {
   280→                            tracing::info!("StartStream received from viewer: {}", start.viewer_id);
   281→                            if let Err(e) = self.init_streaming() {
   282→                                tracing::error!("Failed to init streaming: {}", e);
   283→                            } else {
   284→                                self.state = SessionState::Streaming;
   285→                                self.current_viewer_id = Some(start.viewer_id.clone());
   286→                                tracing::info!("Now streaming to viewer {}", start.viewer_id);
   287→                            }
   288→                            continue;
   289→                        }
   290→                        message::Payload::StopStream(stop) => {
   291→                            tracing::info!("StopStream received for viewer: {}", stop.viewer_id);
   292→                            // Only stop if it matches current viewer
   293→                            if self.current_viewer_id.as_ref() == Some(&stop.viewer_id) {
   294→                                self.release_streaming();
   295→                                self.state = SessionState::Idle;
   296→                                tracing::info!("Stopped streaming, returning to idle mode");
   297→                            }
   298→                            continue;
   299→                        }
   300→                        message::Payload::Heartbeat(hb) => {
   301→                            // Respond to server heartbeat with ack
   302→                            let ack = HeartbeatAck {
   303→                                client_timestamp: hb.timestamp,
   304→                                server_timestamp: chrono::Utc::now().timestamp_millis(),
   305→                            };
   306→                            let ack_msg = Message {
   307→                                payload: Some(message::Payload::HeartbeatAck(ack)),
   308→                            };
   309→                            if let Some(transport) = self.transport.as_mut() {
   310→                                let _ = transport.send(ack_msg).await;
   311→                            }
   312→                            continue;
   313→                        }
   314→                        _ => {}
   315→                    }
   316→                }
   317→
   318→                // Handle other messages (input events, disconnect, etc.)
   319→                self.handle_message(msg).await?;
   320→            }
   321→
   322→            // Check for outgoing chat messages
   323→            if let Some(c) = chat {
   324→                if let Some(outgoing) = c.poll_outgoing() {
   325→                    let chat_proto = ChatMessage {
   326→                        id: outgoing.id,
   327→                        sender: "client".to_string(),
   328→                        content: outgoing.content,
   329→                        timestamp: outgoing.timestamp,
   330→                    };
   331→                    let msg = Message {
   332→                        payload: Some(message::Payload::ChatMessage(chat_proto)),
   333→                    };
   334→                    let transport = self.transport.as_mut().unwrap();
   335→                    transport.send(msg).await?;
   336→                }
   337→            }
   338→
   339→            // State-specific behavior
   340→            match self.state {
   341→                SessionState::Idle => {
   342→                    // In idle mode, just send heartbeats and status periodically
   343→                    if last_heartbeat.elapsed() >= HEARTBEAT_INTERVAL {
   344→                        last_heartbeat = Instant::now();
   345→                        if let Err(e) = self.send_heartbeat().await {
   346→                            tracing::warn!("Failed to send heartbeat: {}", e);
   347→                        }
   348→                    }
   349→
   350→                    if last_status.elapsed() >= STATUS_INTERVAL {
   351→                        last_status = Instant::now();
   352→                        if let Err(e) = self.send_status().await {
   353→                            tracing::warn!("Failed to send status: {}", e);
   354→                        }
   355→                    }
   356→
   357→                    // Periodic update check (only for persistent agents, not support sessions)
   358→                    if self.config.support_code.is_none() && last_update_check.elapsed() >= UPDATE_CHECK_INTERVAL {
   359→                        last_update_check = Instant::now();
   360→                        let server_url = self.config.server_url.replace("/ws/agent", "").replace("wss://", "https://").replace("ws://", "http://");
   361→                        match crate::update::check_for_update(&server_url).await {
   362→                            Ok(Some(version_info)) => {
   363→                                tracing::info!("Update available: {} -> {}", crate::build_info::VERSION, version_info.latest_version);
   364→                                if let Err(e) = crate::update::perform_update(&version_info).await {
   365→                                    tracing::error!("Auto-update failed: {}", e);
   366→                                }
   367→                            }
   368→                            Ok(None) => {
   369→                                tracing::debug!("No update available");
   370→                            }
   371→                            Err(e) => {
   372→                                tracing::debug!("Update check failed: {}", e);
   373→                            }
   374→                        }
   375→                    }
   376→
   377→                    // Longer sleep in idle mode to reduce CPU usage
   378→                    tokio::time::sleep(Duration::from_millis(100)).await;
   379→                }
   380→                SessionState::Streaming => {
   381→                    // In streaming mode, capture and send frames
   382→                    if last_frame_time.elapsed() >= frame_interval {
   383→                        last_frame_time = Instant::now();
   384→
   385→                        if let (Some(capturer), Some(encoder)) =
   386→                            (self.capturer.as_mut(), self.encoder.as_mut())
   387→                        {
   388→                            if let Ok(Some(frame)) = capturer.capture() {
   389→                                if let Ok(encoded) = encoder.encode(&frame) {
   390→                                    if encoded.size > 0 {
   391→                                        let msg = Message {
   392→                                            payload: Some(message::Payload::VideoFrame(encoded.frame)),
   393→                                        };
   394→                                        let transport = self.transport.as_mut().unwrap();
   395→                                        if let Err(e) = transport.send(msg).await {
   396→                                            tracing::warn!("Failed to send frame: {}", e);
   397→                                        }
   398→                                    }
   399→                                }
   400→                            }
   401→                        }
   402→                    }
   403→
   404→                    // Short sleep in streaming mode
   405→                    tokio::time::sleep(Duration::from_millis(1)).await;
   406→                }
   407→                _ => {
   408→                    // Disconnected or connecting - shouldn't be in main loop
   409→                    tokio::time::sleep(Duration::from_millis(100)).await;
   410→                }
   411→            }
   412→
   413→            // Check if still connected
   414→            if let Some(transport) = self.transport.as_ref() {
   415→                if !transport.is_connected() {
   416→                    tracing::warn!("Connection lost");
   417→                    break;
   418→                }
   419→            } else {
   420→                tracing::warn!("Transport is None");
   421→                break;
   422→            }
   423→        }
   424→
   425→        self.release_streaming();
   426→        self.state = SessionState::Disconnected;
   427→        Ok(())
   428→    }
   429→
   430→    /// Handle incoming message from server
   431→    async fn handle_message(&mut self, msg: Message) -> Result<()> {
   432→        match msg.payload {
   433→            Some(message::Payload::MouseEvent(mouse)) => {
   434→                if let Some(input) = self.input.as_mut() {
   435→                    use crate::proto::MouseEventType;
   436→                    use crate::input::MouseButton;
   437→
   438→                    match MouseEventType::try_from(mouse.event_type).unwrap_or(MouseEventType::MouseMove) {
   439→                        MouseEventType::MouseMove => {
   440→                            input.mouse_move(mouse.x, mouse.y)?;
   441→                        }
   442→                        MouseEventType::MouseDown => {
   443→                            input.mouse_move(mouse.x, mouse.y)?;
   444→                            if let Some(ref buttons) = mouse.buttons {
   445→                                if buttons.left { input.mouse_click(MouseButton::Left, true)?; }
   446→                                if buttons.right { input.mouse_click(MouseButton::Right, true)?; }
   447→                                if buttons.middle { input.mouse_click(MouseButton::Middle, true)?; }
   448→                            }
   449→                        }
   450→                        MouseEventType::MouseUp => {
   451→                            if let Some(ref buttons) = mouse.buttons {
   452→                                if buttons.left { input.mouse_click(MouseButton::Left, false)?; }
   453→                                if buttons.right { input.mouse_click(MouseButton::Right, false)?; }
   454→                                if buttons.middle { input.mouse_click(MouseButton::Middle, false)?; }
   455→                            }
   456→                        }
   457→                        MouseEventType::MouseWheel => {
   458→                            input.mouse_scroll(mouse.wheel_delta_x, mouse.wheel_delta_y)?;
   459→                        }
   460→                    }
   461→                }
   462→            }
   463→
   464→            Some(message::Payload::KeyEvent(key)) => {
   465→                if let Some(input) = self.input.as_mut() {
   466→                    input.key_event(key.vk_code as u16, key.down)?;
   467→                }
   468→            }
   469→
   470→            Some(message::Payload::SpecialKey(special)) => {
   471→                if let Some(input) = self.input.as_mut() {
   472→                    use crate::proto::SpecialKey;
   473→                    match SpecialKey::try_from(special.key).ok() {
   474→                        Some(SpecialKey::CtrlAltDel) => {
   475→                            input.send_ctrl_alt_del()?;
   476→                        }
   477→                        _ => {}
   478→                    }
   479→                }
   480→            }
   481→
   482→            Some(message::Payload::AdminCommand(cmd)) => {
   483→                use crate::proto::AdminCommandType;
   484→                tracing::info!("Admin command received: {:?} - {}", cmd.command, cmd.reason);
   485→
   486→                match AdminCommandType::try_from(cmd.command).ok() {
   487→                    Some(AdminCommandType::AdminUninstall) => {
   488→                        tracing::warn!("Uninstall command received from server");
   489→                        // Return special error to trigger uninstall in main loop
   490→                        return Err(anyhow::anyhow!("ADMIN_UNINSTALL: {}", cmd.reason));
   491→                    }
   492→                    Some(AdminCommandType::AdminRestart) => {
   493→                        tracing::info!("Restart command received from server");
   494→                        // For now, just disconnect - the auto-restart logic will handle it
   495→                        return Err(anyhow::anyhow!("ADMIN_RESTART: {}", cmd.reason));
   496→                    }
   497→                    Some(AdminCommandType::AdminUpdate) => {
   498→                        tracing::info!("Update command received from server: {}", cmd.reason);
   499→                        // Trigger update check and perform update if available
   500→                        // The server URL is derived from the config
   501→                        let server_url = self.config.server_url.replace("/ws/agent", "").replace("wss://", "https://").replace("ws://", "http://");
   502→                        match crate::update::check_for_update(&server_url).await {
   503→                            Ok(Some(version_info)) => {
   504→                                tracing::info!("Update available: {} -> {}", crate::build_info::VERSION, version_info.latest_version);
   505→                                if let Err(e) = crate::update::perform_update(&version_info).await {
   506→                                    tracing::error!("Update failed: {}", e);
   507→                                }
   508→                                // If we get here, the update failed (perform_update exits on success)
   509→                            }
   510→                            Ok(None) => {
   511→                                tracing::info!("Already running latest version");
   512→                            }
   513→                            Err(e) => {
   514→                                tracing::error!("Failed to check for updates: {}", e);
   515→                            }
   516→                        }
   517→                    }
   518→                    None => {
   519→                        tracing::warn!("Unknown admin command: {}", cmd.command);
   520→                    }
   521→                }
   522→            }
   523→
   524→            Some(message::Payload::Disconnect(disc)) => {
   525→                tracing::info!("Disconnect requested: {}", disc.reason);
   526→                if disc.reason.contains("cancelled") {
   527→                    return Err(anyhow::anyhow!("SESSION_CANCELLED: {}", disc.reason));
   528→                }
   529→                if disc.reason.contains("administrator") || disc.reason.contains("Disconnected") {
   530→                    return Err(anyhow::anyhow!("ADMIN_DISCONNECT: {}", disc.reason));
   531→                }
   532→                return Err(anyhow::anyhow!("Disconnect: {}", disc.reason));
   533→            }
   534→
   535→            _ => {
   536→                // Ignore unknown messages
   537→            }
   538→        }
   539→
   540→        Ok(())
   541→    }
   542→}
   543→

<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 the agent
//!
//! Handles the lifecycle of a remote session including:
//! - Connection to server
//! - Idle mode (heartbeat only, minimal resources)
//! - Active/streaming mode (capture and send frames)
//! - Input event handling

#[cfg(windows)]
use windows::Win32::System::Console::{AllocConsole, GetConsoleWindow};
#[cfg(windows)]
use windows::Win32::UI::WindowsAndMessaging::{ShowWindow, SW_SHOW};

use crate::capture::{self, Capturer, Display};
use crate::chat::{ChatController, ChatMessage as ChatMsg};
use crate::config::Config;
use crate::encoder::{self, Encoder};
use crate::input::InputController;

/// Show the debug console window (Windows only)
#[cfg(windows)]
fn show_debug_console() {
    unsafe {
        let hwnd = GetConsoleWindow();
        if hwnd.0 == std::ptr::null_mut() {
            let _ = AllocConsole();
            tracing::info!("Debug console window opened");
        } else {
            let _ = ShowWindow(hwnd, SW_SHOW);
            tracing::info!("Debug console window shown");
        }
    }
}

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

use crate::proto::{Message, message, ChatMessage, AgentStatus, Heartbeat, HeartbeatAck};
use crate::transport::WebSocketTransport;
use crate::tray::{TrayController, TrayAction};
use anyhow::Result;
use std::time::{Duration, Instant};

// Heartbeat interval (30 seconds)
const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(30);
// Status report interval (60 seconds)
const STATUS_INTERVAL: Duration = Duration::from_secs(60);
// Update check interval (1 hour)
const UPDATE_CHECK_INTERVAL: Duration = Duration::from_secs(3600);

/// Session manager handles the remote control session
pub struct SessionManager {
    config: Config,
    transport: Option<WebSocketTransport>,
    state: SessionState,
    // Lazy-initialized streaming resources
    capturer: Option<Box<dyn Capturer>>,
    encoder: Option<Box<dyn Encoder>>,
    input: Option<InputController>,
    // Streaming state
    current_viewer_id: Option<String>,
    // System info for status reports
    hostname: String,
    is_elevated: bool,
    start_time: Instant,
}

#[derive(Debug, Clone, PartialEq)]
enum SessionState {
    Disconnected,
    Connecting,
    Idle,       // Connected but not streaming - minimal resource usage
    Streaming,  // Actively capturing and sending frames
}

impl SessionManager {
    /// Create a new session manager
    pub fn new(config: Config, is_elevated: bool) -> Self {
        let hostname = config.hostname();
        Self {
            config,
            transport: None,
            state: SessionState::Disconnected,
            capturer: None,
            encoder: None,
            input: None,
            current_viewer_id: None,
            hostname,
            is_elevated,
            start_time: Instant::now(),
        }
    }

    /// Connect to the server
    pub async fn connect(&mut self) -> Result<()> {
        self.state = SessionState::Connecting;

        let transport = WebSocketTransport::connect(
            &self.config.server_url,
            &self.config.agent_id,
            &self.config.api_key,
            Some(&self.hostname),
            self.config.support_code.as_deref(),
        ).await?;

        self.transport = Some(transport);
        self.state = SessionState::Idle;  // Start in idle mode

        tracing::info!("Connected to server, entering idle mode");

        Ok(())
    }

    /// Initialize streaming resources (capturer, encoder, input)
    fn init_streaming(&mut self) -> Result<()> {
        if self.capturer.is_some() {
            return Ok(()); // Already initialized
        }

        tracing::info!("Initializing streaming resources...");

        // Get primary display
        let primary_display = capture::primary_display()?;
        tracing::info!("Using display: {} ({}x{})",
            primary_display.name, primary_display.width, primary_display.height);

        // Create capturer
        let capturer = capture::create_capturer(
            primary_display.clone(),
            self.config.capture.use_dxgi,
            self.config.capture.gdi_fallback,
        )?;
        self.capturer = Some(capturer);

        // Create encoder
        let encoder = encoder::create_encoder(
            &self.config.encoding.codec,
            self.config.encoding.quality,
        )?;
        self.encoder = Some(encoder);

        // Create input controller
        let input = InputController::new()?;
        self.input = Some(input);

        tracing::info!("Streaming resources initialized");
        Ok(())
    }

    /// Release streaming resources to save CPU/memory when idle
    fn release_streaming(&mut self) {
        if self.capturer.is_some() {
            tracing::info!("Releasing streaming resources");
            self.capturer = None;
            self.encoder = None;
            self.input = None;
            self.current_viewer_id = None;
        }
    }

    /// Get display count for status reports
    fn get_display_count(&self) -> i32 {
        capture::enumerate_displays().map(|d| d.len() as i32).unwrap_or(1)
    }

    /// Send agent status to server
    async fn send_status(&mut self) -> Result<()> {
        let status = AgentStatus {
            hostname: self.hostname.clone(),
            os_version: std::env::consts::OS.to_string(),
            is_elevated: self.is_elevated,
            uptime_secs: self.start_time.elapsed().as_secs() as i64,
            display_count: self.get_display_count(),
            is_streaming: self.state == SessionState::Streaming,
            agent_version: crate::build_info::short_version(),
            organization: self.config.company.clone().unwrap_or_default(),
            site: self.config.site.clone().unwrap_or_default(),
            tags: self.config.tags.clone(),
        };

        let msg = Message {
            payload: Some(message::Payload::AgentStatus(status)),
        };

        if let Some(transport) = self.transport.as_mut() {
            transport.send(msg).await?;
        }

        Ok(())
    }

    /// Send heartbeat to server
    async fn send_heartbeat(&mut self) -> Result<()> {
        let heartbeat = Heartbeat {
            timestamp: chrono::Utc::now().timestamp_millis(),
        };

        let msg = Message {
            payload: Some(message::Payload::Heartbeat(heartbeat)),
        };

        if let Some(transport) = self.transport.as_mut() {
            transport.send(msg).await?;
        }

        Ok(())
    }

    /// Run the session main loop with tray and chat event processing
    pub async fn run_with_tray(&mut self, tray: Option<&TrayController>, chat: Option<&ChatController>) -> Result<()> {
        if self.transport.is_none() {
            anyhow::bail!("Not connected");
        }

        // Send initial status
        self.send_status().await?;

        // Timing for heartbeat and status
        let mut last_heartbeat = Instant::now();
        let mut last_status = Instant::now();
        let mut last_frame_time = Instant::now();
        let mut last_update_check = Instant::now();
        let frame_interval = Duration::from_millis(1000 / self.config.capture.fps as u64);

        // Main loop
        loop {
            // Process tray events
            if let Some(t) = tray {
                if let Some(action) = t.process_events() {
                    match action {
                        TrayAction::EndSession => {
                            tracing::info!("User requested session end via tray");
                            return Err(anyhow::anyhow!("USER_EXIT: Session ended by user"));
                        }
                        TrayAction::ShowDetails => {
                            tracing::info!("User requested details (not yet implemented)");
                        }
                        TrayAction::ShowDebugWindow => {
                            show_debug_console();
                        }
                    }
                }

                if t.exit_requested() {
                    tracing::info!("Exit requested via tray");
                    return Err(anyhow::anyhow!("USER_EXIT: Exit requested by user"));
                }
            }

            // Process incoming messages
            let messages: Vec<Message> = {
                let transport = self.transport.as_mut().unwrap();
                let mut msgs = Vec::new();
                while let Some(msg) = transport.try_recv()? {
                    msgs.push(msg);
                }
                msgs
            };

            for msg in messages {
                // Handle chat messages specially
                if let Some(message::Payload::ChatMessage(chat_msg)) = &msg.payload {
                    if let Some(c) = chat {
                        c.add_message(ChatMsg {
                            id: chat_msg.id.clone(),
                            sender: chat_msg.sender.clone(),
                            content: chat_msg.content.clone(),
                            timestamp: chat_msg.timestamp,
                        });
                    }
                    continue;
                }

                // Handle control messages that affect state
                if let Some(ref payload) = msg.payload {
                    match payload {
                        message::Payload::StartStream(start) => {
                            tracing::info!("StartStream received from viewer: {}", start.viewer_id);
                            if let Err(e) = self.init_streaming() {
                                tracing::error!("Failed to init streaming: {}", e);
                            } else {
                                self.state = SessionState::Streaming;
                                self.current_viewer_id = Some(start.viewer_id.clone());
                                tracing::info!("Now streaming to viewer {}", start.viewer_id);
                            }
                            continue;
                        }
                        message::Payload::StopStream(stop) => {
                            tracing::info!("StopStream received for viewer: {}", stop.viewer_id);
                            // Only stop if it matches current viewer
                            if self.current_viewer_id.as_ref() == Some(&stop.viewer_id) {
                                self.release_streaming();
                                self.state = SessionState::Idle;
                                tracing::info!("Stopped streaming, returning to idle mode");
                            }
                            continue;
                        }
                        message::Payload::Heartbeat(hb) => {
                            // Respond to server heartbeat with ack
                            let ack = HeartbeatAck {
                                client_timestamp: hb.timestamp,
                                server_timestamp: chrono::Utc::now().timestamp_millis(),
                            };
                            let ack_msg = Message {
                                payload: Some(message::Payload::HeartbeatAck(ack)),
                            };
                            if let Some(transport) = self.transport.as_mut() {
                                let _ = transport.send(ack_msg).await;
                            }
                            continue;
                        }
                        _ => {}
                    }
                }

                // Handle other messages (input events, disconnect, etc.)
                self.handle_message(msg).await?;
            }

            // Check for outgoing chat messages
            if let Some(c) = chat {
                if let Some(outgoing) = c.poll_outgoing() {
                    let chat_proto = ChatMessage {
                        id: outgoing.id,
                        sender: "client".to_string(),
                        content: outgoing.content,
                        timestamp: outgoing.timestamp,
                    };
                    let msg = Message {
                        payload: Some(message::Payload::ChatMessage(chat_proto)),
                    };
                    let transport = self.transport.as_mut().unwrap();
                    transport.send(msg).await?;
                }
            }

            // State-specific behavior
            match self.state {
                SessionState::Idle => {
                    // In idle mode, just send heartbeats and status periodically
                    if last_heartbeat.elapsed() >= HEARTBEAT_INTERVAL {
                        last_heartbeat = Instant::now();
                        if let Err(e) = self.send_heartbeat().await {
                            tracing::warn!("Failed to send heartbeat: {}", e);
                        }
                    }

                    if last_status.elapsed() >= STATUS_INTERVAL {
                        last_status = Instant::now();
                        if let Err(e) = self.send_status().await {
                            tracing::warn!("Failed to send status: {}", e);
                        }
                    }

                    // Periodic update check (only for persistent agents, not support sessions)
                    if self.config.support_code.is_none() && last_update_check.elapsed() >= UPDATE_CHECK_INTERVAL {
                        last_update_check = Instant::now();
                        let server_url = self.config.server_url.replace("/ws/agent", "").replace("wss://", "https://").replace("ws://", "http://");
                        match crate::update::check_for_update(&server_url).await {
                            Ok(Some(version_info)) => {
                                tracing::info!("Update available: {} -> {}", crate::build_info::VERSION, version_info.latest_version);
                                if let Err(e) = crate::update::perform_update(&version_info).await {
                                    tracing::error!("Auto-update failed: {}", e);
                                }
                            }
                            Ok(None) => {
                                tracing::debug!("No update available");
                            }
                            Err(e) => {
                                tracing::debug!("Update check failed: {}", e);
                            }
                        }
                    }

                    // Longer sleep in idle mode to reduce CPU usage
                    tokio::time::sleep(Duration::from_millis(100)).await;
                }
                SessionState::Streaming => {
                    // In streaming mode, capture and send frames
                    if last_frame_time.elapsed() >= frame_interval {
                        last_frame_time = Instant::now();

                        if let (Some(capturer), Some(encoder)) =
                            (self.capturer.as_mut(), self.encoder.as_mut())
                        {
                            if let Ok(Some(frame)) = capturer.capture() {
                                if let Ok(encoded) = encoder.encode(&frame) {
                                    if encoded.size > 0 {
                                        let msg = Message {
                                            payload: Some(message::Payload::VideoFrame(encoded.frame)),
                                        };
                                        let transport = self.transport.as_mut().unwrap();
                                        if let Err(e) = transport.send(msg).await {
                                            tracing::warn!("Failed to send frame: {}", e);
                                        }
                                    }
                                }
                            }
                        }
                    }

                    // Short sleep in streaming mode
                    tokio::time::sleep(Duration::from_millis(1)).await;
                }
                _ => {
                    // Disconnected or connecting - shouldn't be in main loop
                    tokio::time::sleep(Duration::from_millis(100)).await;
                }
            }

            // Check if still connected
            if let Some(transport) = self.transport.as_ref() {
                if !transport.is_connected() {
                    tracing::warn!("Connection lost");
                    break;
                }
            } else {
                tracing::warn!("Transport is None");
                break;
            }
        }

        self.release_streaming();
        self.state = SessionState::Disconnected;
        Ok(())
    }

    /// Handle incoming message from server
    async fn handle_message(&mut self, msg: Message) -> Result<()> {
        match msg.payload {
            Some(message::Payload::MouseEvent(mouse)) => {
                if let Some(input) = self.input.as_mut() {
                    use crate::proto::MouseEventType;
                    use crate::input::MouseButton;

                    match MouseEventType::try_from(mouse.event_type).unwrap_or(MouseEventType::MouseMove) {
                        MouseEventType::MouseMove => {
                            input.mouse_move(mouse.x, mouse.y)?;
                        }
                        MouseEventType::MouseDown => {
                            input.mouse_move(mouse.x, mouse.y)?;
                            if let Some(ref buttons) = mouse.buttons {
                                if buttons.left { input.mouse_click(MouseButton::Left, true)?; }
                                if buttons.right { input.mouse_click(MouseButton::Right, true)?; }
                                if buttons.middle { input.mouse_click(MouseButton::Middle, true)?; }
                            }
                        }
                        MouseEventType::MouseUp => {
                            if let Some(ref buttons) = mouse.buttons {
                                if buttons.left { input.mouse_click(MouseButton::Left, false)?; }
                                if buttons.right { input.mouse_click(MouseButton::Right, false)?; }
                                if buttons.middle { input.mouse_click(MouseButton::Middle, false)?; }
                            }
                        }
                        MouseEventType::MouseWheel => {
                            input.mouse_scroll(mouse.wheel_delta_x, mouse.wheel_delta_y)?;
                        }
                    }
                }
            }

            Some(message::Payload::KeyEvent(key)) => {
                if let Some(input) = self.input.as_mut() {
                    input.key_event(key.vk_code as u16, key.down)?;
                }
            }

            Some(message::Payload::SpecialKey(special)) => {
                if let Some(input) = self.input.as_mut() {
                    use crate::proto::SpecialKey;
                    match SpecialKey::try_from(special.key).ok() {
                        Some(SpecialKey::CtrlAltDel) => {
                            input.send_ctrl_alt_del()?;
                        }
                        _ => {}
                    }
                }
            }

            Some(message::Payload::AdminCommand(cmd)) => {
                use crate::proto::AdminCommandType;
                tracing::info!("Admin command received: {:?} - {}", cmd.command, cmd.reason);

                match AdminCommandType::try_from(cmd.command).ok() {
                    Some(AdminCommandType::AdminUninstall) => {
                        tracing::warn!("Uninstall command received from server");
                        // Return special error to trigger uninstall in main loop
                        return Err(anyhow::anyhow!("ADMIN_UNINSTALL: {}", cmd.reason));
                    }
                    Some(AdminCommandType::AdminRestart) => {
                        tracing::info!("Restart command received from server");
                        // For now, just disconnect - the auto-restart logic will handle it
                        return Err(anyhow::anyhow!("ADMIN_RESTART: {}", cmd.reason));
                    }
                    Some(AdminCommandType::AdminUpdate) => {
                        tracing::info!("Update command received from server: {}", cmd.reason);
                        // Trigger update check and perform update if available
                        // The server URL is derived from the config
                        let server_url = self.config.server_url.replace("/ws/agent", "").replace("wss://", "https://").replace("ws://", "http://");
                        match crate::update::check_for_update(&server_url).await {
                            Ok(Some(version_info)) => {
                                tracing::info!("Update available: {} -> {}", crate::build_info::VERSION, version_info.latest_version);
                                if let Err(e) = crate::update::perform_update(&version_info).await {
                                    tracing::error!("Update failed: {}", e);
                                }
                                // If we get here, the update failed (perform_update exits on success)
                            }
                            Ok(None) => {
                                tracing::info!("Already running latest version");
                            }
                            Err(e) => {
                                tracing::error!("Failed to check for updates: {}", e);
                            }
                        }
                    }
                    None => {
                        tracing::warn!("Unknown admin command: {}", cmd.command);
                    }
                }
            }

            Some(message::Payload::Disconnect(disc)) => {
                tracing::info!("Disconnect requested: {}", disc.reason);
                if disc.reason.contains("cancelled") {
                    return Err(anyhow::anyhow!("SESSION_CANCELLED: {}", disc.reason));
                }
                if disc.reason.contains("administrator") || disc.reason.contains("Disconnected") {
                    return Err(anyhow::anyhow!("ADMIN_DISCONNECT: {}", disc.reason));
                }
                return Err(anyhow::anyhow!("Disconnect: {}", disc.reason));
            }

            _ => {
                // Ignore unknown messages
            }
        }

        Ok(())
    }
}


Now let me try the edit again with exactly the same content from the file.