The file C:\Users\MikeSwanson\Claude\guru-connect\agent\src\session\mod.rs has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
     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→
    51→/// Session manager handles the remote control session
    52→pub struct SessionManager {
    53→    config: Config,
    54→    transport: Option<WebSocketTransport>,
    55→    state: SessionState,
    56→    // Lazy-initialized streaming resources
    57→    capturer: Option<Box<dyn Capturer>>,
    58→    encoder: Option<Box<dyn Encoder>>,
    59→    input: Option<InputController>,
    60→    // Streaming state
    61→    current_viewer_id: Option<String>,
    62→    // System info for status reports
    63→    hostname: String,
    64→    is_elevated: bool,
    65→    start_time: Instant,
    66→}
    67→
    68→#[derive(Debug, Clone, PartialEq)]
    69→enum SessionState {
    70→    Disconnected,
    71→    Connecting,
    72→    Idle,       // Connected but not streaming - minimal resource usage
    73→    Streaming,  // Actively capturing and sending frames
    74→}
    75→
    76→impl SessionManager {
    77→    /// Create a new session manager
    78→    pub fn new(config: Config, is_elevated: bool) -> Self {
    79→        let hostname = config.hostname();
    80→        Self {
    81→            config,
    82→            transport: None,
    83→            state: SessionState::Disconnected,
    84→            capturer: None,
    85→            encoder: None,
    86→            input: None,
    87→            current_viewer_id: None,
    88→            hostname,
    89→            is_elevated,
    90→            start_time: Instant::now(),
    91→        }
    92→    }
    93→
    94→    /// Connect to the server
    95→    pub async fn connect(&mut self) -> Result<()> {
    96→        self.state = SessionState::Connecting;
    97→
    98→        let transport = WebSocketTransport::connect(
    99→            &self.config.server_url,
   100→            &self.config.agent_id,
   101→            &self.config.api_key,
   102→            Some(&self.hostname),
   103→            self.config.support_code.as_deref(),
   104→        ).await?;
   105→
   106→        self.transport = Some(transport);
   107→        self.state = SessionState::Idle;  // Start in idle mode
   108→
   109→        tracing::info!("Connected to server, entering idle mode");
   110→
   111→        Ok(())
   112→    }
   113→
   114→    /// Initialize streaming resources (capturer, encoder, input)
   115→    fn init_streaming(&mut self) -> Result<()> {
   116→        if self.capturer.is_some() {
   117→            return Ok(()); // Already initialized
   118→        }
   119→
   120→        tracing::info!("Initializing streaming resources...");
   121→
   122→        // Get primary display
   123→        let primary_display = capture::primary_display()?;
   124→        tracing::info!("Using display: {} ({}x{})",
   125→            primary_display.name, primary_display.width, primary_display.height);
   126→
   127→        // Create capturer
   128→        let capturer = capture::create_capturer(
   129→            primary_display.clone(),
   130→            self.config.capture.use_dxgi,
   131→            self.config.capture.gdi_fallback,
   132→        )?;
   133→        self.capturer = Some(capturer);
   134→
   135→        // Create encoder
   136→        let encoder = encoder::create_encoder(
   137→            &self.config.encoding.codec,
   138→            self.config.encoding.quality,
   139→        )?;
   140→        self.encoder = Some(encoder);
   141→
   142→        // Create input controller
   143→        let input = InputController::new()?;
   144→        self.input = Some(input);
   145→
   146→        tracing::info!("Streaming resources initialized");
   147→        Ok(())
   148→    }
   149→
   150→    /// Release streaming resources to save CPU/memory when idle
   151→    fn release_streaming(&mut self) {
   152→        if self.capturer.is_some() {
   153→            tracing::info!("Releasing streaming resources");
   154→            self.capturer = None;
   155→            self.encoder = None;
   156→            self.input = None;
   157→            self.current_viewer_id = None;
   158→        }
   159→    }
   160→
   161→    /// Get display count for status reports
   162→    fn get_display_count(&self) -> i32 {
   163→        capture::list_displays().map(|d| d.len() as i32).unwrap_or(1)
   164→    }
   165→
   166→    /// Send agent status to server
   167→    async fn send_status(&mut self) -> Result<()> {
   168→        let status = AgentStatus {
   169→            hostname: self.hostname.clone(),
   170→            os_version: std::env::consts::OS.to_string(),
   171→            is_elevated: self.is_elevated,
   172→            uptime_secs: self.start_time.elapsed().as_secs() as i64,
   173→            display_count: self.get_display_count(),
   174→            is_streaming: self.state == SessionState::Streaming,
   175→        };
   176→
   177→        let msg = Message {
   178→            payload: Some(message::Payload::AgentStatus(status)),
   179→        };
   180→
   181→        if let Some(transport) = self.transport.as_mut() {
   182→            transport.send(msg).await?;
   183→        }
   184→
   185→        Ok(())
   186→    }
   187→
   188→    /// Send heartbeat to server
   189→    async fn send_heartbeat(&mut self) -> Result<()> {
   190→        let heartbeat = Heartbeat {
   191→            timestamp: chrono::Utc::now().timestamp_millis(),
   192→        };
   193→
   194→        let msg = Message {
   195→            payload: Some(message::Payload::Heartbeat(heartbeat)),
   196→        };
   197→
   198→        if let Some(transport) = self.transport.as_mut() {
   199→            transport.send(msg).await?;
   200→        }
   201→
   202→        Ok(())
   203→    }
   204→
   205→    /// Run the session main loop with tray and chat event processing
   206→    pub async fn run_with_tray(&mut self, tray: Option<&TrayController>, chat: Option<&ChatController>) -> Result<()> {
   207→        if self.transport.is_none() {
   208→            anyhow::bail!("Not connected");
   209→        }
   210→
   211→        // Send initial status
   212→        self.send_status().await?;
   213→
   214→        // Timing for heartbeat and status
   215→        let mut last_heartbeat = Instant::now();
   216→        let mut last_status = Instant::now();
   217→        let mut last_frame_time = Instant::now();
   218→        let frame_interval = Duration::from_millis(1000 / self.config.capture.fps as u64);
   219→
   220→        // Main loop
   221→        loop {
   222→            // Process tray events
   223→            if let Some(t) = tray {
   224→                if let Some(action) = t.process_events() {
   225→                    match action {
   226→                        TrayAction::EndSession => {
   227→                            tracing::info!("User requested session end via tray");
   228→                            return Err(anyhow::anyhow!("USER_EXIT: Session ended by user"));
   229→                        }
   230→                        TrayAction::ShowDetails => {
   231→                            tracing::info!("User requested details (not yet implemented)");
   232→                        }
   233→                        TrayAction::ShowDebugWindow => {
   234→                            show_debug_console();
   235→                        }
   236→                    }
   237→                }
   238→
   239→                if t.exit_requested() {
   240→                    tracing::info!("Exit requested via tray");
   241→                    return Err(anyhow::anyhow!("USER_EXIT: Exit requested by user"));
   242→                }
   243→            }
   244→
   245→            // Process incoming messages
   246→            let messages: Vec<Message> = {
   247→                let transport = self.transport.as_mut().unwrap();
   248→                let mut msgs = Vec::new();
   249→                while let Some(msg) = transport.try_recv()? {
   250→                    msgs.push(msg);
   251→                }
   252→                msgs
   253→            };
   254→
   255→            for msg in messages {
   256→                // Handle chat messages specially
   257→                if let Some(message::Payload::ChatMessage(chat_msg)) = &msg.payload {
   258→                    if let Some(c) = chat {
   259→                        c.add_message(ChatMsg {
   260→                            id: chat_msg.id.clone(),
   261→                            sender: chat_msg.sender.clone(),
   262→                            content: chat_msg.content.clone(),
   263→                            timestamp: chat_msg.timestamp,
   264→                        });
   265→                    }
   266→                    continue;
   267→                }
   268→
   269→                // Handle control messages that affect state
   270→                if let Some(ref payload) = msg.payload {
   271→                    match payload {
   272→                        message::Payload::StartStream(start) => {
   273→                            tracing::info!("StartStream received from viewer: {}", start.viewer_id);
   274→                            if let Err(e) = self.init_streaming() {
   275→                                tracing::error!("Failed to init streaming: {}", e);
   276→                            } else {
   277→                                self.state = SessionState::Streaming;
   278→                                self.current_viewer_id = Some(start.viewer_id.clone());
   279→                                tracing::info!("Now streaming to viewer {}", start.viewer_id);
   280→                            }
   281→                            continue;
   282→                        }
   283→                        message::Payload::StopStream(stop) => {
   284→                            tracing::info!("StopStream received for viewer: {}", stop.viewer_id);
   285→                            // Only stop if it matches current viewer
   286→                            if self.current_viewer_id.as_ref() == Some(&stop.viewer_id) {
   287→                                self.release_streaming();
   288→                                self.state = SessionState::Idle;
   289→                                tracing::info!("Stopped streaming, returning to idle mode");
   290→                            }
   291→                            continue;
   292→                        }
   293→                        message::Payload::Heartbeat(hb) => {
   294→                            // Respond to server heartbeat with ack
   295→                            let ack = HeartbeatAck {
   296→                                client_timestamp: hb.timestamp,
   297→                                server_timestamp: chrono::Utc::now().timestamp_millis(),
   298→                            };
   299→                            let ack_msg = Message {
   300→                                payload: Some(message::Payload::HeartbeatAck(ack)),
   301→                            };
   302→                            if let Some(transport) = self.transport.as_mut() {
   303→                                let _ = transport.send(ack_msg).await;
   304→                            }
   305→                            continue;
   306→                        }
   307→                        _ => {}
   308→                    }
   309→                }
   310→
   311→                // Handle other messages (input events, disconnect, etc.)
   312→                self.handle_message(msg)?;
   313→            }
   314→
   315→            // Check for outgoing chat messages
   316→            if let Some(c) = chat {
   317→                if let Some(outgoing) = c.poll_outgoing() {
   318→                    let chat_proto = ChatMessage {
   319→                        id: outgoing.id,
   320→                        sender: "client".to_string(),
   321→                        content: outgoing.content,
   322→                        timestamp: outgoing.timestamp,
   323→                    };
   324→                    let msg = Message {
   325→                        payload: Some(message::Payload::ChatMessage(chat_proto)),
   326→                    };
   327→                    let transport = self.transport.as_mut().unwrap();
   328→                    transport.send(msg).await?;
   329→                }
   330→            }
   331→
   332→            // State-specific behavior
   333→            match self.state {
   334→                SessionState::Idle => {
   335→                    // In idle mode, just send heartbeats and status periodically
   336→                    if last_heartbeat.elapsed() >= HEARTBEAT_INTERVAL {
   337→                        last_heartbeat = Instant::now();
   338→                        if let Err(e) = self.send_heartbeat().await {
   339→                            tracing::warn!("Failed to send heartbeat: {}", e);
   340→                        }
   341→                    }
   342→
   343→                    if last_status.elapsed() >= STATUS_INTERVAL {
   344→                        last_status = Instant::now();
   345→                        if let Err(e) = self.send_status().await {
   346→                            tracing::warn!("Failed to send status: {}", e);
   347→                        }
   348→                    }
   349→
   350→                    // Longer sleep in idle mode to reduce CPU usage
   351→                    tokio::time::sleep(Duration::from_millis(100)).await;
   352→                }
   353→                SessionState::Streaming => {
   354→                    // In streaming mode, capture and send frames
   355→                    if last_frame_time.elapsed() >= frame_interval {
   356→                        last_frame_time = Instant::now();
   357→
   358→                        if let (Some(capturer), Some(encoder)) =
   359→                            (self.capturer.as_mut(), self.encoder.as_mut())
   360→                        {
   361→                            if let Ok(Some(frame)) = capturer.capture() {
   362→                                if let Ok(encoded) = encoder.encode(&frame) {
   363→                                    if encoded.size > 0 {
   364→                                        let msg = Message {
   365→                                            payload: Some(message::Payload::VideoFrame(encoded.frame)),
   366→                                        };
   367→                                        let transport = self.transport.as_mut().unwrap();
   368→                                        if let Err(e) = transport.send(msg).await {
   369→                                            tracing::warn!("Failed to send frame: {}", e);
   370→                                        }
   371→                                    }
   372→                                }
   373→                            }
   374→                        }
   375→                    }
   376→
   377→                    // Short sleep in streaming mode
   378→                    tokio::time::sleep(Duration::from_millis(1)).await;
   379→                }
   380→                _ => {
   381→                    // Disconnected or connecting - shouldn't be in main loop
   382→                    tokio::time::sleep(Duration::from_millis(100)).await;
   383→                }
   384→            }
   385→
   386→            // Check if still connected
   387→            if let Some(transport) = self.transport.as_ref() {
   388→                if !transport.is_connected() {
   389→                    tracing::warn!("Connection lost");
   390→                    break;
   391→                }
   392→            } else {
   393→                tracing::warn!("Transport is None");
   394→                break;
   395→            }
   396→        }
   397→
   398→        self.release_streaming();
   399→        self.state = SessionState::Disconnected;
   400→        Ok(())
   401→    }
   402→
   403→    /// Handle incoming message from server
   404→    fn handle_message(&mut self, msg: Message) -> Result<()> {
   405→        match msg.payload {
   406→            Some(message::Payload::MouseEvent(mouse)) => {
   407→                if let Some(input) = self.input.as_mut() {
   408→                    use crate::proto::MouseEventType;
   409→                    use crate::input::MouseButton;
   410→
   411→                    match MouseEventType::try_from(mouse.event_type).unwrap_or(MouseEventType::MouseMove) {
   412→                        MouseEventType::MouseMove => {
   413→                            input.mouse_move(mouse.x, mouse.y)?;
   414→                        }
   415→                        MouseEventType::MouseDown => {
   416→                            input.mouse_move(mouse.x, mouse.y)?;
   417→                            if let Some(ref buttons) = mouse.buttons {
   418→                                if buttons.left { input.mouse_click(MouseButton::Left, true)?; }
   419→                                if buttons.right { input.mouse_click(MouseButton::Right, true)?; }
   420→                                if buttons.middle { input.mouse_click(MouseButton::Middle, true)?; }
   421→                            }
   422→                        }
   423→                        MouseEventType::MouseUp => {
   424→                            if let Some(ref buttons) = mouse.buttons {
   425→                                if buttons.left { input.mouse_click(MouseButton::Left, false)?; }
   426→                                if buttons.right { input.mouse_click(MouseButton::Right, false)?; }
   427→                                if buttons.middle { input.mouse_click(MouseButton::Middle, false)?; }
   428→                            }
   429→                        }
   430→                        MouseEventType::MouseWheel => {
   431→                            input.mouse_scroll(mouse.wheel_delta_x, mouse.wheel_delta_y)?;
   432→                        }
   433→                    }
   434→                }
   435→            }
   436→
   437→            Some(message::Payload::KeyEvent(key)) => {
   438→                if let Some(input) = self.input.as_mut() {
   439→                    input.key_event(key.vk_code as u16, key.down)?;
   440→                }
   441→            }
   442→
   443→            Some(message::Payload::SpecialKey(special)) => {
   444→                if let Some(input) = self.input.as_mut() {
   445→                    use crate::proto::SpecialKey;
   446→                    match SpecialKey::try_from(special.key).ok() {
   447→                        Some(SpecialKey::CtrlAltDel) => {
   448→                            input.send_ctrl_alt_del()?;
   449→                        }
   450→                        _ => {}
   451→                    }
   452→                }
   453→            }
   454→
   455→            Some(message::Payload::Disconnect(disc)) => {
   456→                tracing::info!("Disconnect requested: {}", disc.reason);
   457→                if disc.reason.contains("cancelled") {
   458→                    return Err(anyhow::anyhow!("SESSION_CANCELLED: {}", disc.reason));
   459→                }
   460→                if disc.reason.contains("administrator") || disc.reason.contains("Disconnected") {
   461→                    return Err(anyhow::anyhow!("ADMIN_DISCONNECT: {}", disc.reason));
   462→                }
   463→                return Err(anyhow::anyhow!("Disconnect: {}", disc.reason));
   464→            }
   465→
   466→            _ => {
   467→                // Ignore unknown messages
   468→            }
   469→        }
   470→
   471→        Ok(())
   472→    }
   473→}
   474→

//! 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);

/// 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::list_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,
        };

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

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

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

-//! Handles the lifecycle of a remote session including:

-use crate::capture::{self, Capturer, Display};

-use crate::chat::{ChatController, ChatMessage as ChatMsg};

-/// Show the debug console window (Windows only)

-        if hwnd.0 == std::ptr::null_mut() {

-            tracing::info!("Debug console window opened");

-            let _ = ShowWindow(hwnd, SW_SHOW);

-            tracing::info!("Debug console window shown");

-use crate::proto::{Message, message, ChatMessage};

-/// Session manager handles the remote control session

-    transport: Option<WebSocketTransport>,

-    pub fn new(config: Config) -> Self {

-            state: SessionState::Disconnected,

-    pub async fn connect(&mut self) -> Result<()> {

-        self.state = SessionState::Connecting;

-        let hostname = self.config.hostname();

-        let transport = WebSocketTransport::connect(

-            self.config.support_code.as_deref(),

-        self.transport = Some(transport);

-        self.state = SessionState::Connected;

-    pub async fn run(&mut self) -> Result<()> {

-            anyhow::bail!("Not connected");

-        self.state = SessionState::Active;

-        let primary_display = capture::primary_display()?;

-        tracing::info!("Using display: {} ({}x{})", primary_display.name, primary_display.width, primary_display.height);

-        let mut capturer = capture::create_capturer(

-            self.config.capture.use_dxgi,

-            self.config.capture.gdi_fallback,

-        let mut encoder = encoder::create_encoder(

-            &self.config.encoding.codec,

-            self.config.encoding.quality,

-        let mut input = InputController::new()?;

-        let frame_interval = Duration::from_millis(1000 / self.config.capture.fps as u64);

-        let mut last_frame_time = Instant::now();

-            // Check for incoming messages (non-blocking)

-            // Collect messages first, then release borrow before handling

-            let messages: Vec<Message> = {

-                let transport = self.transport.as_mut().unwrap();

-                let mut msgs = Vec::new();

-                while let Some(msg) = transport.try_recv()? {

-                self.handle_message(&mut input, msg)?;

-            // Capture and send frame if interval elapsed

-            if last_frame_time.elapsed() >= frame_interval {

-                last_frame_time = Instant::now();

-                if let Some(frame) = capturer.capture()? {

-                    let encoded = encoder.encode(&frame)?;

-                    // Skip empty frames (no changes)

-                    if encoded.size > 0 {

-                        let msg = Message {

-                            payload: Some(message::Payload::VideoFrame(encoded.frame)),

-                        let transport = self.transport.as_mut().unwrap();

-                        transport.send(msg).await?;

-            // Small sleep to prevent busy loop

-            tokio::time::sleep(Duration::from_millis(1)).await;

-            // Check if still connected

-            if let Some(transport) = self.transport.as_ref() {

-                if !transport.is_connected() {

-                    tracing::warn!("Connection lost");

-                tracing::warn!("Transport is None");

-        self.state = SessionState::Disconnected;

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

-            anyhow::bail!("Not connected");

-        self.state = SessionState::Active;

-        let primary_display = capture::primary_display()?;

-        let mut capturer = capture::create_capturer(

-            self.config.capture.use_dxgi,

-            self.config.capture.gdi_fallback,

-        let mut encoder = encoder::create_encoder(

-            &self.config.encoding.codec,

-            self.config.encoding.quality,

-        let mut input = InputController::new()?;

-        let mut last_frame_time = Instant::now();

-                if let Some(action) = t.process_events() {

-                        TrayAction::EndSession => {

-                            tracing::info!("User requested session end via tray");

-                            return Err(anyhow::anyhow!("USER_EXIT: Session ended by user"));

-                        TrayAction::ShowDetails => {

-                            // TODO: Show a details dialog

-                            tracing::info!("User requested details (not yet implemented)");

-                        TrayAction::ShowDebugWindow => {

-                            show_debug_console();

-                // Check if exit was requested

-                if t.exit_requested() {

-                    tracing::info!("Exit requested via tray");

-                    return Err(anyhow::anyhow!("USER_EXIT: Exit requested by user"));

-            // Check for incoming messages (non-blocking)

-            let messages: Vec<Message> = {

-                let transport = self.transport.as_mut().unwrap();

-                let mut msgs = Vec::new();

-                while let Some(msg) = transport.try_recv()? {

-                // 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; // Don't pass to handle_message

-                self.handle_message(&mut input, msg)?;

-            // Check for outgoing chat messages from user

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

-            // Capture and send frame if interval elapsed

-            if last_frame_time.elapsed() >= frame_interval {

-                last_frame_time = Instant::now();

-                if let Some(frame) = capturer.capture()? {

-                    let encoded = encoder.encode(&frame)?;

-                    // Skip empty frames (no changes)

-                    if encoded.size > 0 {

-                        let msg = Message {

-                        let transport = self.transport.as_mut().unwrap();

-                        transport.send(msg).await?;

-            // Small sleep to prevent busy loop

-            tokio::time::sleep(Duration::from_millis(1)).await;

-            // Check if still connected

-            if let Some(transport) = self.transport.as_ref() {

-                if !transport.is_connected() {

-                    tracing::warn!("Connection lost");

-                tracing::warn!("Transport is None");

-        self.state = SessionState::Disconnected;

-    /// Handle incoming message from server

-    fn handle_message(&mut self, input: &mut InputController, msg: Message) -> Result<()> {

-            Some(message::Payload::MouseEvent(mouse)) => {

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

-                        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 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)) => {

-                // Handle keyboard event

-                input.key_event(key.vk_code as u16, key.down)?;

-            Some(message::Payload::SpecialKey(special)) => {

-                use crate::proto::SpecialKey;

-                match SpecialKey::try_from(special.key).ok() {

-                    Some(SpecialKey::CtrlAltDel) => {

-                        input.send_ctrl_alt_del()?;

-            Some(message::Payload::Heartbeat(_)) => {

-                // Respond to heartbeat

-                // TODO: Send heartbeat ack

-            Some(message::Payload::Disconnect(disc)) => {

-                tracing::info!("Disconnect requested: {}", disc.reason);

-                // Check if this is a cancellation (support session)

-                if disc.reason.contains("cancelled") {

-                    return Err(anyhow::anyhow!("SESSION_CANCELLED: {}", disc.reason));

-                // Check if this is an admin disconnect (persistent session)

-                if disc.reason.contains("administrator") || disc.reason.contains("Disconnected") {

-                    return Err(anyhow::anyhow!("ADMIN_DISCONNECT: {}", disc.reason));

-                return Err(anyhow::anyhow!("Disconnect: {}", disc.reason));

-                // Ignore unknown messages

+//! Handles the lifecycle of a remote session including:

+//! - Idle mode (heartbeat only, minimal resources)

+//! - Active/streaming mode (capture and send frames)

+use crate::capture::{self, Capturer, Display};

+use crate::chat::{ChatController, ChatMessage as ChatMsg};

+/// Show the debug console window (Windows only)

+        if hwnd.0 == std::ptr::null_mut() {

+            tracing::info!("Debug console window opened");

+            let _ = ShowWindow(hwnd, SW_SHOW);

+            tracing::info!("Debug console window shown");

+use crate::proto::{Message, message, ChatMessage, AgentStatus, Heartbeat, HeartbeatAck};

+const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(30);

+const STATUS_INTERVAL: Duration = Duration::from_secs(60);

+/// Session manager handles the remote control session

+    transport: Option<WebSocketTransport>,

+    // Lazy-initialized streaming resources

+    capturer: Option<Box<dyn Capturer>>,

+    Idle,       // Connected but not streaming - minimal resource usage

+    Streaming,  // Actively capturing and sending frames

+    pub fn new(config: Config, is_elevated: bool) -> Self {

+        let hostname = config.hostname();

+            state: SessionState::Disconnected,

+            start_time: Instant::now(),

+    pub async fn connect(&mut self) -> Result<()> {

+        self.state = SessionState::Connecting;

+        let transport = WebSocketTransport::connect(

+            self.config.support_code.as_deref(),

+        self.transport = Some(transport);

+        self.state = SessionState::Idle;  // Start in idle mode

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

+    /// Initialize streaming resources (capturer, encoder, input)

+    fn init_streaming(&mut self) -> Result<()> {

+            return Ok(()); // Already initialized

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

+        let primary_display = capture::primary_display()?;

+        tracing::info!("Using display: {} ({}x{})",

+            primary_display.name, primary_display.width, primary_display.height);

+        let capturer = capture::create_capturer(

+            self.config.capture.use_dxgi,

+            self.config.capture.gdi_fallback,

+        self.capturer = Some(capturer);

+        let encoder = encoder::create_encoder(

+            &self.config.encoding.codec,

+            self.config.encoding.quality,

+        let input = InputController::new()?;

+        tracing::info!("Streaming resources initialized");

+    /// Release streaming resources to save CPU/memory when idle

+            tracing::info!("Releasing streaming resources");

+            self.current_viewer_id = None;

+    /// Get display count for status reports

+    fn get_display_count(&self) -> i32 {

+        capture::list_displays().map(|d| d.len() as i32).unwrap_or(1)

+    async fn send_status(&mut self) -> Result<()> {

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

+            payload: Some(message::Payload::AgentStatus(status)),

+        if let Some(transport) = self.transport.as_mut() {

+            transport.send(msg).await?;

+    async fn send_heartbeat(&mut self) -> Result<()> {

+            timestamp: chrono::Utc::now().timestamp_millis(),

+            payload: Some(message::Payload::Heartbeat(heartbeat)),

+    /// 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<()> {

+            anyhow::bail!("Not connected");

+        // Timing for heartbeat and status

+        let mut last_heartbeat = Instant::now();

+        let mut last_status = Instant::now();

+        let mut last_frame_time = Instant::now();

+        let frame_interval = Duration::from_millis(1000 / self.config.capture.fps as u64);

+                if let Some(action) = t.process_events() {

+                        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()? {

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

+                // Handle control messages that affect state

+                if let Some(ref payload) = msg.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);

+                                self.state = SessionState::Streaming;

+                                self.current_viewer_id = Some(start.viewer_id.clone());

+                                tracing::info!("Now streaming to viewer {}", start.viewer_id);

+                        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");

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

+                // Handle other messages (input events, disconnect, etc.)

+                self.handle_message(msg)?;

+            // Check for outgoing chat messages

+                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?;

+                SessionState::Idle => {

+                    // In idle mode, just send heartbeats and status periodically

+                    if last_heartbeat.elapsed() >= HEARTBEAT_INTERVAL {

+                        last_heartbeat = Instant::now();

+                        if let Err(e) = self.send_heartbeat().await {

+                            tracing::warn!("Failed to send heartbeat: {}", e);

+                    if last_status.elapsed() >= STATUS_INTERVAL {

+                        last_status = Instant::now();

+                        if let Err(e) = self.send_status().await {

+                            tracing::warn!("Failed to send status: {}", e);

+                    // Longer sleep in idle mode to reduce CPU usage

+                    tokio::time::sleep(Duration::from_millis(100)).await;

+                SessionState::Streaming => {

+                    // In streaming mode, capture and send frames

+                    if last_frame_time.elapsed() >= frame_interval {

+                        last_frame_time = Instant::now();

+                        if let (Some(capturer), Some(encoder)) =

+                            (self.capturer.as_mut(), self.encoder.as_mut())

+                            if let Ok(Some(frame)) = capturer.capture() {

+                                if let Ok(encoded) = encoder.encode(&frame) {

+                                    if encoded.size > 0 {

+                                        let msg = Message {

+                                            payload: Some(message::Payload::VideoFrame(encoded.frame)),

+                                        let transport = self.transport.as_mut().unwrap();

+                                        if let Err(e) = transport.send(msg).await {

+                                            tracing::warn!("Failed to send frame: {}", e);

+                    // Short sleep in streaming mode

+                    tokio::time::sleep(Duration::from_millis(1)).await;

+                    // Disconnected or connecting - shouldn't be in main loop

+                    tokio::time::sleep(Duration::from_millis(100)).await;

+            // Check if still connected

+            if let Some(transport) = self.transport.as_ref() {

+                if !transport.is_connected() {

+                    tracing::warn!("Connection lost");

+                tracing::warn!("Transport is None");

+        self.state = SessionState::Disconnected;

+    /// Handle incoming message from server

+    fn handle_message(&mut self, msg: Message) -> Result<()> {

+            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 => {

+                            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 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)) => {

+                    use crate::proto::SpecialKey;

+                    match SpecialKey::try_from(special.key).ok() {

+                        Some(SpecialKey::CtrlAltDel) => {

+                            input.send_ctrl_alt_del()?;

+            Some(message::Payload::Disconnect(disc)) => {

+                tracing::info!("Disconnect requested: {}", disc.reason);

+                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

//! Session management for the agent
//!
//! Handles the lifecycle of a remote session including:
//! - Connection to server
//! - Authentication
//! - Frame capture and encoding loop
//! - 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};
use crate::transport::WebSocketTransport;
use crate::tray::{TrayController, TrayAction};
use anyhow::Result;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::mpsc;

/// Session manager handles the remote control session
pub struct SessionManager {
    config: Config,
    transport: Option<WebSocketTransport>,
    state: SessionState,
}

#[derive(Debug, Clone, PartialEq)]
enum SessionState {
    Disconnected,
    Connecting,
    Connected,
    Active,
}

impl SessionManager {
    /// Create a new session manager
    pub fn new(config: Config) -> Self {
        Self {
            config,
            transport: None,
            state: SessionState::Disconnected,
        }
    }

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

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

        self.transport = Some(transport);
        self.state = SessionState::Connected;

        Ok(())
    }

    /// Run the session main loop
    pub async fn run(&mut self) -> Result<()> {
        if self.transport.is_none() {
            anyhow::bail!("Not connected");
        }

        self.state = SessionState::Active;

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

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

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

        // Create input controller
        let mut input = InputController::new()?;

        // Calculate frame interval
        let frame_interval = Duration::from_millis(1000 / self.config.capture.fps as u64);
        let mut last_frame_time = Instant::now();

        // Main loop
        loop {
            // Check for incoming messages (non-blocking)
            // Collect messages first, then release borrow before handling
            let messages: Vec<Message> = {
                let transport = self.transport.as_mut().unwrap();
                let mut msgs = Vec::new();
                while let Some(msg) = transport.try_recv()? {
                    msgs.push(msg);
                }
                msgs
            };

            for msg in messages {
                self.handle_message(&mut input, msg)?;
            }

            // Capture and send frame if interval elapsed
            if last_frame_time.elapsed() >= frame_interval {
                last_frame_time = Instant::now();

                if let Some(frame) = capturer.capture()? {
                    let encoded = encoder.encode(&frame)?;

                    // Skip empty frames (no changes)
                    if encoded.size > 0 {
                        let msg = Message {
                            payload: Some(message::Payload::VideoFrame(encoded.frame)),
                        };
                        let transport = self.transport.as_mut().unwrap();
                        transport.send(msg).await?;
                    }
                }
            }

            // Small sleep to prevent busy loop
            tokio::time::sleep(Duration::from_millis(1)).await;

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

        self.state = SessionState::Disconnected;
        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");
        }

        self.state = SessionState::Active;

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

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

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

        // Create input controller
        let mut input = InputController::new()?;

        // Calculate frame interval
        let frame_interval = Duration::from_millis(1000 / self.config.capture.fps as u64);
        let mut last_frame_time = Instant::now();

        // 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 => {
                            // TODO: Show a details dialog
                            tracing::info!("User requested details (not yet implemented)");
                        }
                        TrayAction::ShowDebugWindow => {
                            show_debug_console();
                        }
                    }
                }

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

            // Check for incoming messages (non-blocking)
            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; // Don't pass to handle_message
                }
                self.handle_message(&mut input, msg)?;
            }

            // Check for outgoing chat messages from user
            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?;
                }
            }

            // Capture and send frame if interval elapsed
            if last_frame_time.elapsed() >= frame_interval {
                last_frame_time = Instant::now();

                if let Some(frame) = capturer.capture()? {
                    let encoded = encoder.encode(&frame)?;

                    // Skip empty frames (no changes)
                    if encoded.size > 0 {
                        let msg = Message {
                            payload: Some(message::Payload::VideoFrame(encoded.frame)),
                        };
                        let transport = self.transport.as_mut().unwrap();
                        transport.send(msg).await?;
                    }
                }
            }

            // Small sleep to prevent busy loop
            tokio::time::sleep(Duration::from_millis(1)).await;

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

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

    /// Handle incoming message from server
    fn handle_message(&mut self, input: &mut InputController, msg: Message) -> Result<()> {
        match msg.payload {
            Some(message::Payload::MouseEvent(mouse)) => {
                // Handle mouse event
                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)) => {
                // Handle keyboard event
                input.key_event(key.vk_code as u16, key.down)?;
            }

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

            Some(message::Payload::Heartbeat(_)) => {
                // Respond to heartbeat
                // TODO: Send heartbeat ack
            }

            Some(message::Payload::Disconnect(disc)) => {
                tracing::info!("Disconnect requested: {}", disc.reason);
                // Check if this is a cancellation (support session)
                if disc.reason.contains("cancelled") {
                    return Err(anyhow::anyhow!("SESSION_CANCELLED: {}", disc.reason));
                }
                // Check if this is an admin disconnect (persistent session)
                if disc.reason.contains("administrator") || disc.reason.contains("Disconnected") {
                    return Err(anyhow::anyhow!("ADMIN_DISCONNECT: {}", disc.reason));
                }
                return Err(anyhow::anyhow!("Disconnect: {}", disc.reason));
            }

            _ => {
                // Ignore unknown messages
            }
        }

        Ok(())
    }
}


Good, now I need to update main.rs to pass is_elevated to SessionManager::new(). Let me also update the todos and then look at main.rs.


Now I need to update main.rs to pass `is_elevated` to SessionManager:


TodoWrite

Add StartStream/StopStream/Heartbeat messages to protobuf

Implement agent state machine (idle vs streaming)

Add heartbeat mechanism with timeout detection

Update server to route viewer connections and send start/stop

Build and deploy updated server and agent


Todos have been modified successfully. Ensure that you continue to use the todo list to track your progress. Please proceed with the current tasks if applicable

Add StartStream/StopStream/Heartbeat messages to protobuf

Implement agent state machine (idle vs streaming)

Add heartbeat mechanism with timeout detection

Update server to route viewer connections and send start/stop

Build and deploy updated server and agent