     1→//! Session management for the agent
     2→//!
     3→//! Handles the lifecycle of a remote session including:
     4→//! - Connection to server
     5→//! - Authentication
     6→//! - Frame capture and encoding loop
     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→use crate::proto::{Message, message, ChatMessage};
    40→use crate::transport::WebSocketTransport;
    41→use crate::tray::{TrayController, TrayAction};
    42→use anyhow::Result;
    43→use std::sync::Arc;
    44→use std::time::{Duration, Instant};
    45→use tokio::sync::mpsc;
    46→
    47→/// Session manager handles the remote control session
    48→pub struct SessionManager {
    49→    config: Config,
    50→    transport: Option<WebSocketTransport>,
    51→    state: SessionState,
    52→}
    53→
    54→#[derive(Debug, Clone, PartialEq)]
    55→enum SessionState {
    56→    Disconnected,
    57→    Connecting,
    58→    Connected,
    59→    Active,
    60→}
    61→
    62→impl SessionManager {
    63→    /// Create a new session manager
    64→    pub fn new(config: Config) -> Self {
    65→        Self {
    66→            config,
    67→            transport: None,
    68→            state: SessionState::Disconnected,
    69→        }
    70→    }
    71→
    72→    /// Connect to the server
    73→    pub async fn connect(&mut self) -> Result<()> {
    74→        self.state = SessionState::Connecting;
    75→
    76→        let hostname = self.config.hostname();
    77→        let transport = WebSocketTransport::connect(
    78→            &self.config.server_url,
    79→            &self.config.agent_id,
    80→            &self.config.api_key,
    81→            Some(&hostname),
    82→            self.config.support_code.as_deref(),
    83→        ).await?;
    84→
    85→        self.transport = Some(transport);
    86→        self.state = SessionState::Connected;
    87→
    88→        Ok(())
    89→    }
    90→
    91→    /// Run the session main loop
    92→    pub async fn run(&mut self) -> Result<()> {
    93→        if self.transport.is_none() {
    94→            anyhow::bail!("Not connected");
    95→        }
    96→
    97→        self.state = SessionState::Active;
    98→
    99→        // Get primary display
   100→        let primary_display = capture::primary_display()?;
   101→        tracing::info!("Using display: {} ({}x{})", primary_display.name, primary_display.width, primary_display.height);
   102→
   103→        // Create capturer
   104→        let mut capturer = capture::create_capturer(
   105→            primary_display.clone(),
   106→            self.config.capture.use_dxgi,
   107→            self.config.capture.gdi_fallback,
   108→        )?;
   109→
   110→        // Create encoder
   111→        let mut encoder = encoder::create_encoder(
   112→            &self.config.encoding.codec,
   113→            self.config.encoding.quality,
   114→        )?;
   115→
   116→        // Create input controller
   117→        let mut input = InputController::new()?;
   118→
   119→        // Calculate frame interval
   120→        let frame_interval = Duration::from_millis(1000 / self.config.capture.fps as u64);
   121→        let mut last_frame_time = Instant::now();
   122→
   123→        // Main loop
   124→        loop {
   125→            // Check for incoming messages (non-blocking)
   126→            // Collect messages first, then release borrow before handling
   127→            let messages: Vec<Message> = {
   128→                let transport = self.transport.as_mut().unwrap();
   129→                let mut msgs = Vec::new();
   130→                while let Some(msg) = transport.try_recv()? {
   131→                    msgs.push(msg);
   132→                }
   133→                msgs
   134→            };
   135→
   136→            for msg in messages {
   137→                self.handle_message(&mut input, msg)?;
   138→            }
   139→
   140→            // Capture and send frame if interval elapsed
   141→            if last_frame_time.elapsed() >= frame_interval {
   142→                last_frame_time = Instant::now();
   143→
   144→                if let Some(frame) = capturer.capture()? {
   145→                    let encoded = encoder.encode(&frame)?;
   146→
   147→                    // Skip empty frames (no changes)
   148→                    if encoded.size > 0 {
   149→                        let msg = Message {
   150→                            payload: Some(message::Payload::VideoFrame(encoded.frame)),
   151→                        };
   152→                        let transport = self.transport.as_mut().unwrap();
   153→                        transport.send(msg).await?;
   154→                    }
   155→                }
   156→            }
   157→
   158→            // Small sleep to prevent busy loop
   159→            tokio::time::sleep(Duration::from_millis(1)).await;
   160→
   161→            // Check if still connected
   162→            if let Some(transport) = self.transport.as_ref() {
   163→                if !transport.is_connected() {
   164→                    tracing::warn!("Connection lost");
   165→                    break;
   166→                }
   167→            } else {
   168→                tracing::warn!("Transport is None");
   169→                break;
   170→            }
   171→        }
   172→
   173→        self.state = SessionState::Disconnected;
   174→        Ok(())
   175→    }
   176→
   177→    /// Run the session main loop with tray and chat event processing
   178→    pub async fn run_with_tray(&mut self, tray: Option<&TrayController>, chat: Option<&ChatController>) -> Result<()> {
   179→        if self.transport.is_none() {
   180→            anyhow::bail!("Not connected");
   181→        }
   182→
   183→        self.state = SessionState::Active;
   184→
   185→        // Get primary display
   186→        let primary_display = capture::primary_display()?;
   187→        tracing::info!("Using display: {} ({}x{})", primary_display.name, primary_display.width, primary_display.height);
   188→
   189→        // Create capturer
   190→        let mut capturer = capture::create_capturer(
   191→            primary_display.clone(),
   192→            self.config.capture.use_dxgi,
   193→            self.config.capture.gdi_fallback,
   194→        )?;
   195→
   196→        // Create encoder
   197→        let mut encoder = encoder::create_encoder(
   198→            &self.config.encoding.codec,
   199→            self.config.encoding.quality,
   200→        )?;
   201→
   202→        // Create input controller
   203→        let mut input = InputController::new()?;
   204→
   205→        // Calculate frame interval
   206→        let frame_interval = Duration::from_millis(1000 / self.config.capture.fps as u64);
   207→        let mut last_frame_time = Instant::now();
   208→
   209→        // Main loop
   210→        loop {
   211→            // Process tray events
   212→            if let Some(t) = tray {
   213→                if let Some(action) = t.process_events() {
   214→                    match action {
   215→                        TrayAction::EndSession => {
   216→                            tracing::info!("User requested session end via tray");
   217→                            return Err(anyhow::anyhow!("USER_EXIT: Session ended by user"));
   218→                        }
   219→                        TrayAction::ShowDetails => {
   220→                            // TODO: Show a details dialog
   221→                            tracing::info!("User requested details (not yet implemented)");
   222→                        }
   223→                        TrayAction::ShowDebugWindow => {
   224→                            show_debug_console();
   225→                        }
   226→                    }
   227→                }
   228→
   229→                // Check if exit was requested
   230→                if t.exit_requested() {
   231→                    tracing::info!("Exit requested via tray");
   232→                    return Err(anyhow::anyhow!("USER_EXIT: Exit requested by user"));
   233→                }
   234→            }
   235→
   236→            // Check for incoming messages (non-blocking)
   237→            let messages: Vec<Message> = {
   238→                let transport = self.transport.as_mut().unwrap();
   239→                let mut msgs = Vec::new();
   240→                while let Some(msg) = transport.try_recv()? {
   241→                    msgs.push(msg);
   242→                }
   243→                msgs
   244→            };
   245→
   246→            for msg in messages {
   247→                // Handle chat messages specially
   248→                if let Some(message::Payload::ChatMessage(chat_msg)) = &msg.payload {
   249→                    if let Some(c) = chat {
   250→                        c.add_message(ChatMsg {
   251→                            id: chat_msg.id.clone(),
   252→                            sender: chat_msg.sender.clone(),
   253→                            content: chat_msg.content.clone(),
   254→                            timestamp: chat_msg.timestamp,
   255→                        });
   256→                    }
   257→                    continue; // Don't pass to handle_message
   258→                }
   259→                self.handle_message(&mut input, msg)?;
   260→            }
   261→
   262→            // Check for outgoing chat messages from user
   263→            if let Some(c) = chat {
   264→                if let Some(outgoing) = c.poll_outgoing() {
   265→                    let chat_proto = ChatMessage {
   266→                        id: outgoing.id,
   267→                        sender: "client".to_string(),
   268→                        content: outgoing.content,
   269→                        timestamp: outgoing.timestamp,
   270→                    };
   271→                    let msg = Message {
   272→                        payload: Some(message::Payload::ChatMessage(chat_proto)),
   273→                    };
   274→                    let transport = self.transport.as_mut().unwrap();
   275→                    transport.send(msg).await?;
   276→                }
   277→            }
   278→
   279→            // Capture and send frame if interval elapsed
   280→            if last_frame_time.elapsed() >= frame_interval {
   281→                last_frame_time = Instant::now();
   282→
   283→                if let Some(frame) = capturer.capture()? {
   284→                    let encoded = encoder.encode(&frame)?;
   285→
   286→                    // Skip empty frames (no changes)
   287→                    if encoded.size > 0 {
   288→                        let msg = Message {
   289→                            payload: Some(message::Payload::VideoFrame(encoded.frame)),
   290→                        };
   291→                        let transport = self.transport.as_mut().unwrap();
   292→                        transport.send(msg).await?;
   293→                    }
   294→                }
   295→            }
   296→
   297→            // Small sleep to prevent busy loop
   298→            tokio::time::sleep(Duration::from_millis(1)).await;
   299→
   300→            // Check if still connected
   301→            if let Some(transport) = self.transport.as_ref() {
   302→                if !transport.is_connected() {
   303→                    tracing::warn!("Connection lost");
   304→                    break;
   305→                }
   306→            } else {
   307→                tracing::warn!("Transport is None");
   308→                break;
   309→            }
   310→        }
   311→
   312→        self.state = SessionState::Disconnected;
   313→        Ok(())
   314→    }
   315→
   316→    /// Handle incoming message from server
   317→    fn handle_message(&mut self, input: &mut InputController, msg: Message) -> Result<()> {
   318→        match msg.payload {
   319→            Some(message::Payload::MouseEvent(mouse)) => {
   320→                // Handle mouse event
   321→                use crate::proto::MouseEventType;
   322→                use crate::input::MouseButton;
   323→
   324→                match MouseEventType::try_from(mouse.event_type).unwrap_or(MouseEventType::MouseMove) {
   325→                    MouseEventType::MouseMove => {
   326→                        input.mouse_move(mouse.x, mouse.y)?;
   327→                    }
   328→                    MouseEventType::MouseDown => {
   329→                        input.mouse_move(mouse.x, mouse.y)?;
   330→                        if let Some(ref buttons) = mouse.buttons {
   331→                            if buttons.left { input.mouse_click(MouseButton::Left, true)?; }
   332→                            if buttons.right { input.mouse_click(MouseButton::Right, true)?; }
   333→                            if buttons.middle { input.mouse_click(MouseButton::Middle, true)?; }
   334→                        }
   335→                    }
   336→                    MouseEventType::MouseUp => {
   337→                        if let Some(ref buttons) = mouse.buttons {
   338→                            if buttons.left { input.mouse_click(MouseButton::Left, false)?; }
   339→                            if buttons.right { input.mouse_click(MouseButton::Right, false)?; }
   340→                            if buttons.middle { input.mouse_click(MouseButton::Middle, false)?; }
   341→                        }
   342→                    }
   343→                    MouseEventType::MouseWheel => {
   344→                        input.mouse_scroll(mouse.wheel_delta_x, mouse.wheel_delta_y)?;
   345→                    }
   346→                }
   347→            }
   348→
   349→            Some(message::Payload::KeyEvent(key)) => {
   350→                // Handle keyboard event
   351→                input.key_event(key.vk_code as u16, key.down)?;
   352→            }
   353→
   354→            Some(message::Payload::SpecialKey(special)) => {
   355→                use crate::proto::SpecialKey;
   356→                match SpecialKey::try_from(special.key).ok() {
   357→                    Some(SpecialKey::CtrlAltDel) => {
   358→                        input.send_ctrl_alt_del()?;
   359→                    }
   360→                    _ => {}
   361→                }
   362→            }
   363→
   364→            Some(message::Payload::Heartbeat(_)) => {
   365→                // Respond to heartbeat
   366→                // TODO: Send heartbeat ack
   367→            }
   368→
   369→            Some(message::Payload::Disconnect(disc)) => {
   370→                tracing::info!("Disconnect requested: {}", disc.reason);
   371→                // Check if this is a cancellation (support session)
   372→                if disc.reason.contains("cancelled") {
   373→                    return Err(anyhow::anyhow!("SESSION_CANCELLED: {}", disc.reason));
   374→                }
   375→                // Check if this is an admin disconnect (persistent session)
   376→                if disc.reason.contains("administrator") || disc.reason.contains("Disconnected") {
   377→                    return Err(anyhow::anyhow!("ADMIN_DISCONNECT: {}", disc.reason));
   378→                }
   379→                return Err(anyhow::anyhow!("Disconnect: {}", disc.reason));
   380→            }
   381→
   382→            _ => {
   383→                // Ignore unknown messages
   384→            }
   385→        }
   386→
   387→        Ok(())
   388→    }
   389→}
   390→

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


Now I need to rewrite the session module to implement idle/active mode. The key changes:

1. Change SessionState to include Idle and Streaming states
2. Create lazy-initialized capturer and encoder (only when streaming)
3. Main loop behavior depends on state:
   - Idle: Send heartbeats, handle messages (including StartStream)
   - Streaming: Capture frames, send video, handle StopStream
4. Send AgentStatus periodically when idle

Let me rewrite the session module.


Now I'll rewrite the session module with idle/active state machine:


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

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