Write

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

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


<tool_use_error>File has not been read yet. Read it first before writing to it.</tool_use_error>

Error: File has not been read yet. Read it first before writing to it.


Hmm, it says the file hasn't been read. Let me re-read it before writing.


     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 Write tool again now that the file has been read.


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

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


There's something wrong with the file read tracking. Let me try a different approach - I'll use the Edit tool to edit a specific section of the file instead of Write.


Let me try using Edit instead: