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

//! WebSocket relay handlers
//!
//! Handles WebSocket connections from agents and viewers,
//! relaying video frames and input events between them.

use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        Query, State,
    },
    response::IntoResponse,
};
use futures_util::{SinkExt, StreamExt};
use prost::Message as ProstMessage;
use serde::Deserialize;
use tracing::{error, info, warn};
use uuid::Uuid;

use crate::proto;
use crate::session::SessionManager;
use crate::AppState;

#[derive(Debug, Deserialize)]
pub struct AgentParams {
    agent_id: String,
    #[serde(default)]
    agent_name: Option<String>,
    #[serde(default)]
    support_code: Option<String>,
    #[serde(default)]
    hostname: Option<String>,
}

#[derive(Debug, Deserialize)]
pub struct ViewerParams {
    session_id: String,
}

/// WebSocket handler for agent connections
pub async fn agent_ws_handler(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    Query(params): Query<AgentParams>,
) -> impl IntoResponse {
    let agent_id = params.agent_id;
    let agent_name = params.hostname.or(params.agent_name).unwrap_or_else(|| agent_id.clone());
    let support_code = params.support_code;
    let sessions = state.sessions.clone();
    let support_codes = state.support_codes.clone();

    ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, agent_id, agent_name, support_code))
}

/// WebSocket handler for viewer connections
pub async fn viewer_ws_handler(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    Query(params): Query<ViewerParams>,
) -> impl IntoResponse {
    let session_id = params.session_id;
    let sessions = state.sessions.clone();

    ws.on_upgrade(move |socket| handle_viewer_connection(socket, sessions, session_id))
}

/// Handle an agent WebSocket connection
async fn handle_agent_connection(
    socket: WebSocket,
    sessions: SessionManager,
    support_codes: crate::support_codes::SupportCodeManager,
    agent_id: String,
    agent_name: String,
    support_code: Option<String>,
) {
    info!("Agent connected: {} ({})", agent_name, agent_id);

    let (mut ws_sender, mut ws_receiver) = socket.split();

    // If a support code was provided, check if it's valid
    if let Some(ref code) = support_code {
        // Check if the code is cancelled or invalid
        if support_codes.is_cancelled(code).await {
            warn!("Agent tried to connect with cancelled code: {}", code);
            // Send disconnect message to agent
            let disconnect_msg = proto::Message {
                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
                    reason: "Support session was cancelled by technician".to_string(),
                })),
            };
            let mut buf = Vec::new();
            if prost::Message::encode(&disconnect_msg, &mut buf).is_ok() {
                let _ = ws_sender.send(Message::Binary(buf.into())).await;
            }
            let _ = ws_sender.close().await;
            return;
        }
    }

    // Register the agent and get channels
    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone()).await;

    info!("Session created: {} (agent in idle mode)", session_id);

    // If a support code was provided, mark it as connected
    if let Some(ref code) = support_code {
        info!("Linking support code {} to session {}", code, session_id);
        support_codes.mark_connected(code, Some(agent_name.clone()), Some(agent_id.clone())).await;
        support_codes.link_session(code, session_id).await;
    }

    // Use Arc<Mutex> for sender so we can use it from multiple places
    let ws_sender = std::sync::Arc::new(tokio::sync::Mutex::new(ws_sender));
    let ws_sender_input = ws_sender.clone();
    let ws_sender_cancel = ws_sender.clone();

    // Task to forward input events from viewers to agent
    let input_forward = tokio::spawn(async move {
        while let Some(input_data) = input_rx.recv().await {
            let mut sender = ws_sender_input.lock().await;
            if sender.send(Message::Binary(input_data.into())).await.is_err() {
                break;
            }
        }
    });

    let sessions_cleanup = sessions.clone();
    let sessions_status = sessions.clone();
    let support_codes_cleanup = support_codes.clone();
    let support_code_cleanup = support_code.clone();
    let support_code_check = support_code.clone();
    let support_codes_check = support_codes.clone();

    // Task to check for cancellation every 2 seconds
    let cancel_check = tokio::spawn(async move {
        let mut interval = tokio::time::interval(std::time::Duration::from_secs(2));
        loop {
            interval.tick().await;
            if let Some(ref code) = support_code_check {
                if support_codes_check.is_cancelled(code).await {
                    info!("Support code {} was cancelled, disconnecting agent", code);
                    // Send disconnect message
                    let disconnect_msg = proto::Message {
                        payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
                            reason: "Support session was cancelled by technician".to_string(),
                        })),
                    };
                    let mut buf = Vec::new();
                    if prost::Message::encode(&disconnect_msg, &mut buf).is_ok() {
                        let mut sender = ws_sender_cancel.lock().await;
                        let _ = sender.send(Message::Binary(buf.into())).await;
                        let _ = sender.close().await;
                    }
                    break;
                }
            }
        }
    });

    // Main loop: receive messages from agent
    while let Some(msg) = ws_receiver.next().await {
        match msg {
            Ok(Message::Binary(data)) => {
                // Try to decode as protobuf message
                match proto::Message::decode(data.as_ref()) {
                    Ok(proto_msg) => {
                        match &proto_msg.payload {
                            Some(proto::message::Payload::VideoFrame(_)) => {
                                // Broadcast frame to all viewers (only sent when streaming)
                                let _ = frame_tx.send(data.to_vec());
                            }
                            Some(proto::message::Payload::ChatMessage(chat)) => {
                                // Broadcast chat message to all viewers
                                info!("Chat from client: {}", chat.content);
                                let _ = frame_tx.send(data.to_vec());
                            }
                            Some(proto::message::Payload::AgentStatus(status)) => {
                                // Update session with agent status
                                sessions_status.update_agent_status(
                                    session_id,
                                    Some(status.os_version.clone()),
                                    status.is_elevated,
                                    status.uptime_secs,
                                    status.display_count,
                                    status.is_streaming,
                                ).await;
                                info!("Agent status update: {} - streaming={}, uptime={}s",
                                    status.hostname, status.is_streaming, status.uptime_secs);
                            }
                            Some(proto::message::Payload::Heartbeat(_)) => {
                                // Update heartbeat timestamp
                                sessions_status.update_heartbeat(session_id).await;
                            }
                            Some(proto::message::Payload::HeartbeatAck(_)) => {
                                // Agent acknowledged our heartbeat
                                sessions_status.update_heartbeat(session_id).await;
                            }
                            _ => {}
                        }
                    }
                    Err(e) => {
                        warn!("Failed to decode agent message: {}", e);
                    }
                }
            }
            Ok(Message::Close(_)) => {
                info!("Agent disconnected: {}", agent_id);
                break;
            }
            Ok(Message::Ping(data)) => {
                // Pong is handled automatically by axum
                let _ = data;
            }
            Ok(_) => {}
            Err(e) => {
                error!("WebSocket error from agent {}: {}", agent_id, e);
                break;
            }
        }
    }

    // Cleanup
    input_forward.abort();
    cancel_check.abort();
    sessions_cleanup.remove_session(session_id).await;

    // Mark support code as completed if one was used (unless cancelled)
    if let Some(ref code) = support_code_cleanup {
        if !support_codes_cleanup.is_cancelled(code).await {
            support_codes_cleanup.mark_completed(code).await;
            info!("Support code {} marked as completed", code);
        }
    }

    info!("Session {} ended", session_id);
}

/// Handle a viewer WebSocket connection
async fn handle_viewer_connection(
    socket: WebSocket,
    sessions: SessionManager,
    session_id_str: String,
) {
    // Parse session ID
    let session_id = match uuid::Uuid::parse_str(&session_id_str) {
        Ok(id) => id,
        Err(_) => {
            warn!("Invalid session ID: {}", session_id_str);
            return;
        }
    };

    // Generate unique viewer ID
    let viewer_id = Uuid::new_v4().to_string();

    // Join the session (this sends StartStream to agent if first viewer)
    let (mut frame_rx, input_tx) = match sessions.join_session(session_id, viewer_id.clone()).await {
        Some(channels) => channels,
        None => {
            warn!("Session not found: {}", session_id);
            return;
        }
    };

    info!("Viewer {} joined session: {}", viewer_id, session_id);

    let (mut ws_sender, mut ws_receiver) = socket.split();

    // Task to forward frames from agent to this viewer
    let frame_forward = tokio::spawn(async move {
        while let Ok(frame_data) = frame_rx.recv().await {
            if ws_sender.send(Message::Binary(frame_data.into())).await.is_err() {
                break;
            }
        }
    });

    let sessions_cleanup = sessions.clone();
    let viewer_id_cleanup = viewer_id.clone();

    // Main loop: receive input from viewer and forward to agent
    while let Some(msg) = ws_receiver.next().await {
        match msg {
            Ok(Message::Binary(data)) => {
                // Try to decode as protobuf message
                match proto::Message::decode(data.as_ref()) {
                    Ok(proto_msg) => {
                        match &proto_msg.payload {
                            Some(proto::message::Payload::MouseEvent(_)) |
                            Some(proto::message::Payload::KeyEvent(_)) |
                            Some(proto::message::Payload::SpecialKey(_)) => {
                                // Forward input to agent
                                let _ = input_tx.send(data.to_vec()).await;
                            }
                            Some(proto::message::Payload::ChatMessage(chat)) => {
                                // Forward chat message to agent
                                info!("Chat from technician: {}", chat.content);
                                let _ = input_tx.send(data.to_vec()).await;
                            }
                            _ => {}
                        }
                    }
                    Err(e) => {
                        warn!("Failed to decode viewer message: {}", e);
                    }
                }
            }
            Ok(Message::Close(_)) => {
                info!("Viewer {} disconnected from session: {}", viewer_id, session_id);
                break;
            }
            Ok(_) => {}
            Err(e) => {
                error!("WebSocket error from viewer {}: {}", viewer_id, e);
                break;
            }
        }
    }

    // Cleanup (this sends StopStream to agent if last viewer)
    frame_forward.abort();
    sessions_cleanup.leave_session(session_id, &viewer_id_cleanup).await;
    info!("Viewer {} left session: {}", viewer_id_cleanup, session_id);
}

-//! Handles WebSocket connections from agents and viewers,

-//! relaying video frames and input events between them.

-        ws::{Message, WebSocket, WebSocketUpgrade},

-/// WebSocket handler for agent connections

-    let agent_name = params.hostname.or(params.agent_name).unwrap_or_else(|| agent_id.clone());

-    let support_code = params.support_code;

-    let sessions = state.sessions.clone();

-    let support_codes = state.support_codes.clone();

-    ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, agent_id, agent_name, support_code))

-/// WebSocket handler for viewer connections

-    Query(params): Query<ViewerParams>,

-    let session_id = params.session_id;

-    ws.on_upgrade(move |socket| handle_viewer_connection(socket, sessions, session_id))

-/// Handle an agent WebSocket connection

-    support_codes: crate::support_codes::SupportCodeManager,

-    info!("Agent connected: {} ({})", agent_name, agent_id);

-    let (mut ws_sender, mut ws_receiver) = socket.split();

-    // If a support code was provided, check if it's valid

-    if let Some(ref code) = support_code {

-        // Check if the code is cancelled or invalid

-        if support_codes.is_cancelled(code).await {

-            warn!("Agent tried to connect with cancelled code: {}", code);

-            // Send disconnect message to agent

-            let disconnect_msg = proto::Message {

-                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {

-                    reason: "Support session was cancelled by technician".to_string(),

-            if prost::Message::encode(&disconnect_msg, &mut buf).is_ok() {

-                let _ = ws_sender.send(Message::Binary(buf.into())).await;

-            let _ = ws_sender.close().await;

-    // Register the agent and get channels

-    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone()).await;

-    info!("Session created: {}", session_id);

-    // If a support code was provided, mark it as connected

-        info!("Linking support code {} to session {}", code, session_id);

-        support_codes.mark_connected(code, Some(agent_name.clone()), Some(agent_id.clone())).await;

-        support_codes.link_session(code, session_id).await;

-    // Use Arc<Mutex> for sender so we can use it from multiple places

-    let ws_sender = std::sync::Arc::new(tokio::sync::Mutex::new(ws_sender));

-    let ws_sender_input = ws_sender.clone();

-    let ws_sender_cancel = ws_sender.clone();

-    // Task to forward input events from viewers to agent

-    let input_forward = tokio::spawn(async move {

-        while let Some(input_data) = input_rx.recv().await {

-            let mut sender = ws_sender_input.lock().await;

-            if sender.send(Message::Binary(input_data.into())).await.is_err() {

-    let sessions_cleanup = sessions.clone();

-    let support_codes_cleanup = support_codes.clone();

-    let support_code_cleanup = support_code.clone();

-    let support_code_check = support_code.clone();

-    let support_codes_check = support_codes.clone();

-    // Task to check for cancellation every 2 seconds

-    let cancel_check = tokio::spawn(async move {

-        let mut interval = tokio::time::interval(std::time::Duration::from_secs(2));

-            if let Some(ref code) = support_code_check {

-                if support_codes_check.is_cancelled(code).await {

-                    info!("Support code {} was cancelled, disconnecting agent", code);

-                    // Send disconnect message

-                    let disconnect_msg = proto::Message {

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

-                        let mut sender = ws_sender_cancel.lock().await;

-                        let _ = sender.send(Message::Binary(buf.into())).await;

-                        let _ = sender.close().await;

-    // Main loop: receive frames from agent and broadcast to viewers

-    while let Some(msg) = ws_receiver.next().await {

-            Ok(Message::Binary(data)) => {

-                // Try to decode as protobuf message

-                match proto::Message::decode(data.as_ref()) {

-                        match &proto_msg.payload {

-                            Some(proto::message::Payload::VideoFrame(_)) => {

-                                // Broadcast frame to all viewers

-                                let _ = frame_tx.send(data.to_vec());

-                            Some(proto::message::Payload::ChatMessage(chat)) => {

-                                // Broadcast chat message to all viewers

-                                info!("Chat from client: {}", chat.content);

-                        warn!("Failed to decode agent message: {}", e);

-                info!("Agent disconnected: {}", agent_id);

-            Ok(Message::Ping(data)) => {

-                // Pong is handled automatically by axum

-                error!("WebSocket error from agent {}: {}", agent_id, e);

-    sessions_cleanup.remove_session(session_id).await;

-    // Mark support code as completed if one was used (unless cancelled)

-    if let Some(ref code) = support_code_cleanup {

-        if !support_codes_cleanup.is_cancelled(code).await {

-            support_codes_cleanup.mark_completed(code).await;

-            info!("Support code {} marked as completed", code);

-    info!("Session {} ended", session_id);

-/// Handle a viewer WebSocket connection

-    let session_id = match uuid::Uuid::parse_str(&session_id_str) {

-            warn!("Invalid session ID: {}", session_id_str);

-    let (mut frame_rx, input_tx) = match sessions.join_session(session_id).await {

-            warn!("Session not found: {}", session_id);

-    info!("Viewer joined session: {}", session_id);

-    let (mut ws_sender, mut ws_receiver) = socket.split();

-    // Task to forward frames from agent to this viewer

-    let frame_forward = tokio::spawn(async move {

-        while let Ok(frame_data) = frame_rx.recv().await {

-            if ws_sender.send(Message::Binary(frame_data.into())).await.is_err() {

-    let sessions_cleanup = sessions.clone();

-    // Main loop: receive input from viewer and forward to agent

-    while let Some(msg) = ws_receiver.next().await {

-            Ok(Message::Binary(data)) => {

-                // Try to decode as protobuf message

-                match proto::Message::decode(data.as_ref()) {

-                        match &proto_msg.payload {

-                            Some(proto::message::Payload::MouseEvent(_)) |

-                            Some(proto::message::Payload::KeyEvent(_)) => {

-                                // Forward input to agent

-                                let _ = input_tx.send(data.to_vec()).await;

-                            Some(proto::message::Payload::ChatMessage(chat)) => {

-                                // Forward chat message to agent

-                                info!("Chat from technician: {}", chat.content);

-                        warn!("Failed to decode viewer message: {}", e);

-                info!("Viewer disconnected from session: {}", session_id);

-                error!("WebSocket error from viewer: {}", e);

-    sessions_cleanup.leave_session(session_id).await;

-    info!("Viewer left session: {}", session_id);

+//! Handles WebSocket connections from agents and viewers,

+//! relaying video frames and input events between them.

+        ws::{Message, WebSocket, WebSocketUpgrade},

+/// WebSocket handler for agent connections

+    let agent_name = params.hostname.or(params.agent_name).unwrap_or_else(|| agent_id.clone());

+    let support_code = params.support_code;

+    let sessions = state.sessions.clone();

+    let support_codes = state.support_codes.clone();

+    ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, agent_id, agent_name, support_code))

+/// WebSocket handler for viewer connections

+    Query(params): Query<ViewerParams>,

+    let session_id = params.session_id;

+    ws.on_upgrade(move |socket| handle_viewer_connection(socket, sessions, session_id))

+/// Handle an agent WebSocket connection

+    support_codes: crate::support_codes::SupportCodeManager,

+    info!("Agent connected: {} ({})", agent_name, agent_id);

+    let (mut ws_sender, mut ws_receiver) = socket.split();

+    // If a support code was provided, check if it's valid

+    if let Some(ref code) = support_code {

+        // Check if the code is cancelled or invalid

+        if support_codes.is_cancelled(code).await {

+            warn!("Agent tried to connect with cancelled code: {}", code);

+            // Send disconnect message to agent

+            let disconnect_msg = proto::Message {

+                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {

+                    reason: "Support session was cancelled by technician".to_string(),

+            if prost::Message::encode(&disconnect_msg, &mut buf).is_ok() {

+                let _ = ws_sender.send(Message::Binary(buf.into())).await;

+            let _ = ws_sender.close().await;

+    // Register the agent and get channels

+    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone()).await;

+    info!("Session created: {} (agent in idle mode)", session_id);

+    // If a support code was provided, mark it as connected

+        info!("Linking support code {} to session {}", code, session_id);

+        support_codes.mark_connected(code, Some(agent_name.clone()), Some(agent_id.clone())).await;

+        support_codes.link_session(code, session_id).await;

+    // Use Arc<Mutex> for sender so we can use it from multiple places

+    let ws_sender = std::sync::Arc::new(tokio::sync::Mutex::new(ws_sender));

+    let ws_sender_input = ws_sender.clone();

+    let ws_sender_cancel = ws_sender.clone();

+    // Task to forward input events from viewers to agent

+    let input_forward = tokio::spawn(async move {

+        while let Some(input_data) = input_rx.recv().await {

+            let mut sender = ws_sender_input.lock().await;

+            if sender.send(Message::Binary(input_data.into())).await.is_err() {

+    let sessions_cleanup = sessions.clone();

+    let sessions_status = sessions.clone();

+    let support_codes_cleanup = support_codes.clone();

+    let support_code_cleanup = support_code.clone();

+    let support_code_check = support_code.clone();

+    let support_codes_check = support_codes.clone();

+    // Task to check for cancellation every 2 seconds

+    let cancel_check = tokio::spawn(async move {

+        let mut interval = tokio::time::interval(std::time::Duration::from_secs(2));

+            if let Some(ref code) = support_code_check {

+                if support_codes_check.is_cancelled(code).await {

+                    info!("Support code {} was cancelled, disconnecting agent", code);

+                    // Send disconnect message

+                    let disconnect_msg = proto::Message {

+                    let mut buf = Vec::new();

+                        let mut sender = ws_sender_cancel.lock().await;

+                        let _ = sender.send(Message::Binary(buf.into())).await;

+                        let _ = sender.close().await;

+    // Main loop: receive messages from agent

+    while let Some(msg) = ws_receiver.next().await {

+            Ok(Message::Binary(data)) => {

+                // Try to decode as protobuf message

+                match proto::Message::decode(data.as_ref()) {

+                        match &proto_msg.payload {

+                            Some(proto::message::Payload::VideoFrame(_)) => {

+                                // Broadcast frame to all viewers (only sent when streaming)

+                                let _ = frame_tx.send(data.to_vec());

+                            Some(proto::message::Payload::ChatMessage(chat)) => {

+                                // Broadcast chat message to all viewers

+                                info!("Chat from client: {}", chat.content);

+                            Some(proto::message::Payload::AgentStatus(status)) => {

+                                // Update session with agent status

+                                sessions_status.update_agent_status(

+                                    session_id,

+                                    Some(status.os_version.clone()),

+                                    status.is_elevated,

+                                    status.uptime_secs,

+                                    status.display_count,

+                                    status.is_streaming,

+                                ).await;

+                                info!("Agent status update: {} - streaming={}, uptime={}s",

+                                    status.hostname, status.is_streaming, status.uptime_secs);

+                            Some(proto::message::Payload::Heartbeat(_)) => {

+                                // Update heartbeat timestamp

+                                sessions_status.update_heartbeat(session_id).await;

+                            Some(proto::message::Payload::HeartbeatAck(_)) => {

+                                // Agent acknowledged our heartbeat

+                        warn!("Failed to decode agent message: {}", e);

+                info!("Agent disconnected: {}", agent_id);

+            Ok(Message::Ping(data)) => {

+                // Pong is handled automatically by axum

+                error!("WebSocket error from agent {}: {}", agent_id, e);

+    sessions_cleanup.remove_session(session_id).await;

+    // Mark support code as completed if one was used (unless cancelled)

+    if let Some(ref code) = support_code_cleanup {

+        if !support_codes_cleanup.is_cancelled(code).await {

+            support_codes_cleanup.mark_completed(code).await;

+            info!("Support code {} marked as completed", code);

+    info!("Session {} ended", session_id);

+/// Handle a viewer WebSocket connection

+    let session_id = match uuid::Uuid::parse_str(&session_id_str) {

+            warn!("Invalid session ID: {}", session_id_str);

+    let viewer_id = Uuid::new_v4().to_string();

+    // Join the session (this sends StartStream to agent if first viewer)

+    let (mut frame_rx, input_tx) = match sessions.join_session(session_id, viewer_id.clone()).await {

+            warn!("Session not found: {}", session_id);

+    info!("Viewer {} joined session: {}", viewer_id, session_id);

+    let (mut ws_sender, mut ws_receiver) = socket.split();

+    // Task to forward frames from agent to this viewer

+    let frame_forward = tokio::spawn(async move {

+        while let Ok(frame_data) = frame_rx.recv().await {

+            if ws_sender.send(Message::Binary(frame_data.into())).await.is_err() {

+    let sessions_cleanup = sessions.clone();

+    let viewer_id_cleanup = viewer_id.clone();

+    // Main loop: receive input from viewer and forward to agent

+    while let Some(msg) = ws_receiver.next().await {

+            Ok(Message::Binary(data)) => {

+                // Try to decode as protobuf message

+                match proto::Message::decode(data.as_ref()) {

+                        match &proto_msg.payload {

+                            Some(proto::message::Payload::MouseEvent(_)) |

+                            Some(proto::message::Payload::KeyEvent(_)) |

+                            Some(proto::message::Payload::SpecialKey(_)) => {

+                                // Forward input to agent

+                                let _ = input_tx.send(data.to_vec()).await;

+                            Some(proto::message::Payload::ChatMessage(chat)) => {

+                                // Forward chat message to agent

+                                info!("Chat from technician: {}", chat.content);

+                        warn!("Failed to decode viewer message: {}", e);

+                info!("Viewer {} disconnected from session: {}", viewer_id, session_id);

+                error!("WebSocket error from viewer {}: {}", viewer_id, e);

+    // Cleanup (this sends StopStream to agent if last viewer)

+    sessions_cleanup.leave_session(session_id, &viewer_id_cleanup).await;

+    info!("Viewer {} left session: {}", viewer_id_cleanup, session_id);

//! WebSocket relay handlers
//!
//! Handles WebSocket connections from agents and viewers,
//! relaying video frames and input events between them.

use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        Query, State,
    },
    response::IntoResponse,
};
use futures_util::{SinkExt, StreamExt};
use prost::Message as ProstMessage;
use serde::Deserialize;
use tracing::{error, info, warn};

use crate::proto;
use crate::session::SessionManager;
use crate::AppState;

#[derive(Debug, Deserialize)]
pub struct AgentParams {
    agent_id: String,
    #[serde(default)]
    agent_name: Option<String>,
    #[serde(default)]
    support_code: Option<String>,
    #[serde(default)]
    hostname: Option<String>,
}

#[derive(Debug, Deserialize)]
pub struct ViewerParams {
    session_id: String,
}

/// WebSocket handler for agent connections
pub async fn agent_ws_handler(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    Query(params): Query<AgentParams>,
) -> impl IntoResponse {
    let agent_id = params.agent_id;
    let agent_name = params.hostname.or(params.agent_name).unwrap_or_else(|| agent_id.clone());
    let support_code = params.support_code;
    let sessions = state.sessions.clone();
    let support_codes = state.support_codes.clone();

    ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, agent_id, agent_name, support_code))
}

/// WebSocket handler for viewer connections
pub async fn viewer_ws_handler(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    Query(params): Query<ViewerParams>,
) -> impl IntoResponse {
    let session_id = params.session_id;
    let sessions = state.sessions.clone();

    ws.on_upgrade(move |socket| handle_viewer_connection(socket, sessions, session_id))
}

/// Handle an agent WebSocket connection
async fn handle_agent_connection(
    socket: WebSocket,
    sessions: SessionManager,
    support_codes: crate::support_codes::SupportCodeManager,
    agent_id: String,
    agent_name: String,
    support_code: Option<String>,
) {
    info!("Agent connected: {} ({})", agent_name, agent_id);

    let (mut ws_sender, mut ws_receiver) = socket.split();

    // If a support code was provided, check if it's valid
    if let Some(ref code) = support_code {
        // Check if the code is cancelled or invalid
        if support_codes.is_cancelled(code).await {
            warn!("Agent tried to connect with cancelled code: {}", code);
            // Send disconnect message to agent
            let disconnect_msg = proto::Message {
                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
                    reason: "Support session was cancelled by technician".to_string(),
                })),
            };
            let mut buf = Vec::new();
            if prost::Message::encode(&disconnect_msg, &mut buf).is_ok() {
                let _ = ws_sender.send(Message::Binary(buf.into())).await;
            }
            let _ = ws_sender.close().await;
            return;
        }
    }

    // Register the agent and get channels
    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone()).await;

    info!("Session created: {}", session_id);

    // If a support code was provided, mark it as connected
    if let Some(ref code) = support_code {
        info!("Linking support code {} to session {}", code, session_id);
        support_codes.mark_connected(code, Some(agent_name.clone()), Some(agent_id.clone())).await;
        support_codes.link_session(code, session_id).await;
    }

    // Use Arc<Mutex> for sender so we can use it from multiple places
    let ws_sender = std::sync::Arc::new(tokio::sync::Mutex::new(ws_sender));
    let ws_sender_input = ws_sender.clone();
    let ws_sender_cancel = ws_sender.clone();

    // Task to forward input events from viewers to agent
    let input_forward = tokio::spawn(async move {
        while let Some(input_data) = input_rx.recv().await {
            let mut sender = ws_sender_input.lock().await;
            if sender.send(Message::Binary(input_data.into())).await.is_err() {
                break;
            }
        }
    });

    let sessions_cleanup = sessions.clone();
    let support_codes_cleanup = support_codes.clone();
    let support_code_cleanup = support_code.clone();
    let support_code_check = support_code.clone();
    let support_codes_check = support_codes.clone();

    // Task to check for cancellation every 2 seconds
    let cancel_check = tokio::spawn(async move {
        let mut interval = tokio::time::interval(std::time::Duration::from_secs(2));
        loop {
            interval.tick().await;
            if let Some(ref code) = support_code_check {
                if support_codes_check.is_cancelled(code).await {
                    info!("Support code {} was cancelled, disconnecting agent", code);
                    // Send disconnect message
                    let disconnect_msg = proto::Message {
                        payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
                            reason: "Support session was cancelled by technician".to_string(),
                        })),
                    };
                    let mut buf = Vec::new();
                    if prost::Message::encode(&disconnect_msg, &mut buf).is_ok() {
                        let mut sender = ws_sender_cancel.lock().await;
                        let _ = sender.send(Message::Binary(buf.into())).await;
                        let _ = sender.close().await;
                    }
                    break;
                }
            }
        }
    });

    // Main loop: receive frames from agent and broadcast to viewers
    while let Some(msg) = ws_receiver.next().await {
        match msg {
            Ok(Message::Binary(data)) => {
                // Try to decode as protobuf message
                match proto::Message::decode(data.as_ref()) {
                    Ok(proto_msg) => {
                        match &proto_msg.payload {
                            Some(proto::message::Payload::VideoFrame(_)) => {
                                // Broadcast frame to all viewers
                                let _ = frame_tx.send(data.to_vec());
                            }
                            Some(proto::message::Payload::ChatMessage(chat)) => {
                                // Broadcast chat message to all viewers
                                info!("Chat from client: {}", chat.content);
                                let _ = frame_tx.send(data.to_vec());
                            }
                            _ => {}
                        }
                    }
                    Err(e) => {
                        warn!("Failed to decode agent message: {}", e);
                    }
                }
            }
            Ok(Message::Close(_)) => {
                info!("Agent disconnected: {}", agent_id);
                break;
            }
            Ok(Message::Ping(data)) => {
                // Pong is handled automatically by axum
                let _ = data;
            }
            Ok(_) => {}
            Err(e) => {
                error!("WebSocket error from agent {}: {}", agent_id, e);
                break;
            }
        }
    }

    // Cleanup
    input_forward.abort();
    cancel_check.abort();
    sessions_cleanup.remove_session(session_id).await;

    // Mark support code as completed if one was used (unless cancelled)
    if let Some(ref code) = support_code_cleanup {
        if !support_codes_cleanup.is_cancelled(code).await {
            support_codes_cleanup.mark_completed(code).await;
            info!("Support code {} marked as completed", code);
        }
    }

    info!("Session {} ended", session_id);
}

/// Handle a viewer WebSocket connection
async fn handle_viewer_connection(
    socket: WebSocket,
    sessions: SessionManager,
    session_id_str: String,
) {
    // Parse session ID
    let session_id = match uuid::Uuid::parse_str(&session_id_str) {
        Ok(id) => id,
        Err(_) => {
            warn!("Invalid session ID: {}", session_id_str);
            return;
        }
    };

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

    info!("Viewer joined session: {}", session_id);

    let (mut ws_sender, mut ws_receiver) = socket.split();

    // Task to forward frames from agent to this viewer
    let frame_forward = tokio::spawn(async move {
        while let Ok(frame_data) = frame_rx.recv().await {
            if ws_sender.send(Message::Binary(frame_data.into())).await.is_err() {
                break;
            }
        }
    });

    let sessions_cleanup = sessions.clone();

    // Main loop: receive input from viewer and forward to agent
    while let Some(msg) = ws_receiver.next().await {
        match msg {
            Ok(Message::Binary(data)) => {
                // Try to decode as protobuf message
                match proto::Message::decode(data.as_ref()) {
                    Ok(proto_msg) => {
                        match &proto_msg.payload {
                            Some(proto::message::Payload::MouseEvent(_)) |
                            Some(proto::message::Payload::KeyEvent(_)) => {
                                // Forward input to agent
                                let _ = input_tx.send(data.to_vec()).await;
                            }
                            Some(proto::message::Payload::ChatMessage(chat)) => {
                                // Forward chat message to agent
                                info!("Chat from technician: {}", chat.content);
                                let _ = input_tx.send(data.to_vec()).await;
                            }
                            _ => {}
                        }
                    }
                    Err(e) => {
                        warn!("Failed to decode viewer message: {}", e);
                    }
                }
            }
            Ok(Message::Close(_)) => {
                info!("Viewer disconnected from session: {}", session_id);
                break;
            }
            Ok(_) => {}
            Err(e) => {
                error!("WebSocket error from viewer: {}", e);
                break;
            }
        }
    }

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


Now I need to check if there's a `list_displays` function in the capture module. Let me check and add it if missing. Also I should update the API response to include the new session fields like is_streaming and last_heartbeat.

Let me check the capture module for list_displays and the main.rs for the API routes.


Now I need to check if `list_displays` exists in the capture module and update the API response:


     1→//! GuruConnect Server - WebSocket Relay Server
     2→//!
     3→//! Handles connections from both agents and dashboard viewers,
     4→//! relaying video frames and input events between them.
     5→
     6→mod config;
     7→mod relay;
     8→mod session;
     9→mod auth;
    10→mod api;
    11→mod db;
    12→mod support_codes;
    13→
    14→pub mod proto {
    15→    include!(concat!(env!("OUT_DIR"), "/guruconnect.rs"));
    16→}
    17→
    18→use anyhow::Result;
    19→use axum::{
    20→    Router,
    21→    routing::{get, post},
    22→    extract::{Path, State, Json},
    23→    response::{Html, IntoResponse},
    24→    http::StatusCode,
    25→};
    26→use std::net::SocketAddr;
    27→use tower_http::cors::{Any, CorsLayer};
    28→use tower_http::trace::TraceLayer;
    29→use tower_http::services::ServeDir;
    30→use tracing::{info, Level};
    31→use tracing_subscriber::FmtSubscriber;
    32→use serde::Deserialize;
    33→
    34→use support_codes::{SupportCodeManager, CreateCodeRequest, SupportCode, CodeValidation};
    35→
    36→/// Application state
    37→#[derive(Clone)]
    38→pub struct AppState {
    39→    sessions: session::SessionManager,
    40→    support_codes: SupportCodeManager,
    41→}
    42→
    43→#[tokio::main]
    44→async fn main() -> Result<()> {
    45→    // Initialize logging
    46→    let _subscriber = FmtSubscriber::builder()
    47→        .with_max_level(Level::INFO)
    48→        .with_target(true)
    49→        .init();
    50→
    51→    info!("GuruConnect Server v{}", env!("CARGO_PKG_VERSION"));
    52→
    53→    // Load configuration
    54→    let config = config::Config::load()?;
    55→
    56→    // Use port 3002 for GuruConnect
    57→    let listen_addr = std::env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:3002".to_string());
    58→    info!("Loaded configuration, listening on {}", listen_addr);
    59→
    60→    // Create application state
    61→    let state = AppState {
    62→        sessions: session::SessionManager::new(),
    63→        support_codes: SupportCodeManager::new(),
    64→    };
    65→
    66→    // Build router
    67→    let app = Router::new()
    68→        // Health check
    69→        .route("/health", get(health))
    70→
    71→        // Portal API - Support codes
    72→        .route("/api/codes", post(create_code))
    73→        .route("/api/codes", get(list_codes))
    74→        .route("/api/codes/:code/validate", get(validate_code))
    75→        .route("/api/codes/:code/cancel", post(cancel_code))
    76→
    77→        // WebSocket endpoints
    78→        .route("/ws/agent", get(relay::agent_ws_handler))
    79→        .route("/ws/viewer", get(relay::viewer_ws_handler))
    80→
    81→        // REST API - Sessions
    82→        .route("/api/sessions", get(list_sessions))
    83→        .route("/api/sessions/:id", get(get_session))
    84→        .route("/api/sessions/:id", axum::routing::delete(disconnect_session))
    85→
    86→        // HTML page routes (clean URLs)
    87→        .route("/login", get(serve_login))
    88→        .route("/dashboard", get(serve_dashboard))
    89→
    90→        // State
    91→        .with_state(state)
    92→
    93→        // Serve static files for portal (fallback)
    94→        .fallback_service(ServeDir::new("static").append_index_html_on_directories(true))
    95→
    96→        // Middleware
    97→        .layer(TraceLayer::new_for_http())
    98→        .layer(
    99→            CorsLayer::new()
   100→                .allow_origin(Any)
   101→                .allow_methods(Any)
   102→                .allow_headers(Any),
   103→        );
   104→
   105→    // Start server
   106→    let addr: SocketAddr = listen_addr.parse()?;
   107→    let listener = tokio::net::TcpListener::bind(addr).await?;
   108→
   109→    info!("Server listening on {}", addr);
   110→
   111→    axum::serve(listener, app).await?;
   112→
   113→    Ok(())
   114→}
   115→
   116→async fn health() -> &'static str {
   117→    "OK"
   118→}
   119→
   120→// Support code API handlers
   121→
   122→async fn create_code(
   123→    State(state): State<AppState>,
   124→    Json(request): Json<CreateCodeRequest>,
   125→) -> Json<SupportCode> {
   126→    let code = state.support_codes.create_code(request).await;
   127→    info!("Created support code: {}", code.code);
   128→    Json(code)
   129→}
   130→
   131→async fn list_codes(
   132→    State(state): State<AppState>,
   133→) -> Json<Vec<SupportCode>> {
   134→    Json(state.support_codes.list_active_codes().await)
   135→}
   136→
   137→#[derive(Deserialize)]
   138→struct ValidateParams {
   139→    code: String,
   140→}
   141→
   142→async fn validate_code(
   143→    State(state): State<AppState>,
   144→    Path(code): Path<String>,
   145→) -> Json<CodeValidation> {
   146→    Json(state.support_codes.validate_code(&code).await)
   147→}
   148→
   149→async fn cancel_code(
   150→    State(state): State<AppState>,
   151→    Path(code): Path<String>,
   152→) -> impl IntoResponse {
   153→    if state.support_codes.cancel_code(&code).await {
   154→        (StatusCode::OK, "Code cancelled")
   155→    } else {
   156→        (StatusCode::BAD_REQUEST, "Cannot cancel code")
   157→    }
   158→}
   159→
   160→// Session API handlers (updated to use AppState)
   161→
   162→async fn list_sessions(
   163→    State(state): State<AppState>,
   164→) -> Json<Vec<api::SessionInfo>> {
   165→    let sessions = state.sessions.list_sessions().await;
   166→    Json(sessions.into_iter().map(api::SessionInfo::from).collect())
   167→}
   168→
   169→async fn get_session(
   170→    State(state): State<AppState>,
   171→    Path(id): Path<String>,
   172→) -> Result<Json<api::SessionInfo>, (StatusCode, &'static str)> {
   173→    let session_id = uuid::Uuid::parse_str(&id)
   174→        .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID"))?;
   175→
   176→    let session = state.sessions.get_session(session_id).await
   177→        .ok_or((StatusCode::NOT_FOUND, "Session not found"))?;
   178→
   179→    Ok(Json(api::SessionInfo::from(session)))
   180→}
   181→
   182→async fn disconnect_session(
   183→    State(state): State<AppState>,
   184→    Path(id): Path<String>,
   185→) -> impl IntoResponse {
   186→    let session_id = match uuid::Uuid::parse_str(&id) {
   187→        Ok(id) => id,
   188→        Err(_) => return (StatusCode::BAD_REQUEST, "Invalid session ID"),
   189→    };
   190→
   191→    if state.sessions.disconnect_session(session_id, "Disconnected by administrator").await {
   192→        info!("Session {} disconnected by admin", session_id);
   193→        (StatusCode::OK, "Session disconnected")
   194→    } else {
   195→        (StatusCode::NOT_FOUND, "Session not found")
   196→    }
   197→}
   198→
   199→// Static page handlers
   200→async fn serve_login() -> impl IntoResponse {
   201→    match tokio::fs::read_to_string("static/login.html").await {
   202→        Ok(content) => Html(content).into_response(),
   203→        Err(_) => (StatusCode::NOT_FOUND, "Page not found").into_response(),
   204→    }
   205→}
   206→
   207→async fn serve_dashboard() -> impl IntoResponse {
   208→    match tokio::fs::read_to_string("static/dashboard.html").await {
   209→        Ok(content) => Html(content).into_response(),
   210→        Err(_) => (StatusCode::NOT_FOUND, "Page not found").into_response(),
   211→    }
   212→}
   213→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

//! GuruConnect Server - WebSocket Relay Server
//!
//! Handles connections from both agents and dashboard viewers,
//! relaying video frames and input events between them.

mod config;
mod relay;
mod session;
mod auth;
mod api;
mod db;
mod support_codes;

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

use anyhow::Result;
use axum::{
    Router,
    routing::{get, post},
    extract::{Path, State, Json},
    response::{Html, IntoResponse},
    http::StatusCode,
};
use std::net::SocketAddr;
use tower_http::cors::{Any, CorsLayer};
use tower_http::trace::TraceLayer;
use tower_http::services::ServeDir;
use tracing::{info, Level};
use tracing_subscriber::FmtSubscriber;
use serde::Deserialize;

use support_codes::{SupportCodeManager, CreateCodeRequest, SupportCode, CodeValidation};

/// Application state
#[derive(Clone)]
pub struct AppState {
    sessions: session::SessionManager,
    support_codes: SupportCodeManager,
}

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

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

    // Load configuration
    let config = config::Config::load()?;

    // Use port 3002 for GuruConnect
    let listen_addr = std::env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:3002".to_string());
    info!("Loaded configuration, listening on {}", listen_addr);

    // Create application state
    let state = AppState {
        sessions: session::SessionManager::new(),
        support_codes: SupportCodeManager::new(),
    };

    // Build router
    let app = Router::new()
        // Health check
        .route("/health", get(health))

        // Portal API - Support codes
        .route("/api/codes", post(create_code))
        .route("/api/codes", get(list_codes))
        .route("/api/codes/:code/validate", get(validate_code))
        .route("/api/codes/:code/cancel", post(cancel_code))

        // WebSocket endpoints
        .route("/ws/agent", get(relay::agent_ws_handler))
        .route("/ws/viewer", get(relay::viewer_ws_handler))

        // REST API - Sessions
        .route("/api/sessions", get(list_sessions))
        .route("/api/sessions/:id", get(get_session))
        .route("/api/sessions/:id", axum::routing::delete(disconnect_session))

        // HTML page routes (clean URLs)
        .route("/login", get(serve_login))
        .route("/dashboard", get(serve_dashboard))

        // State
        .with_state(state)

        // Serve static files for portal (fallback)
        .fallback_service(ServeDir::new("static").append_index_html_on_directories(true))

        // Middleware
        .layer(TraceLayer::new_for_http())
        .layer(
            CorsLayer::new()
                .allow_origin(Any)
                .allow_methods(Any)
                .allow_headers(Any),
        );

    // Start server
    let addr: SocketAddr = listen_addr.parse()?;
    let listener = tokio::net::TcpListener::bind(addr).await?;

    info!("Server listening on {}", addr);

    axum::serve(listener, app).await?;

    Ok(())
}

async fn health() -> &'static str {
    "OK"
}

// Support code API handlers

async fn create_code(
    State(state): State<AppState>,
    Json(request): Json<CreateCodeRequest>,
) -> Json<SupportCode> {
    let code = state.support_codes.create_code(request).await;
    info!("Created support code: {}", code.code);
    Json(code)
}

async fn list_codes(
    State(state): State<AppState>,
) -> Json<Vec<SupportCode>> {
    Json(state.support_codes.list_active_codes().await)
}

#[derive(Deserialize)]
struct ValidateParams {
    code: String,
}

async fn validate_code(
    State(state): State<AppState>,
    Path(code): Path<String>,
) -> Json<CodeValidation> {
    Json(state.support_codes.validate_code(&code).await)
}

async fn cancel_code(
    State(state): State<AppState>,
    Path(code): Path<String>,
) -> impl IntoResponse {
    if state.support_codes.cancel_code(&code).await {
        (StatusCode::OK, "Code cancelled")
    } else {
        (StatusCode::BAD_REQUEST, "Cannot cancel code")
    }
}

// Session API handlers (updated to use AppState)

async fn list_sessions(
    State(state): State<AppState>,
) -> Json<Vec<api::SessionInfo>> {
    let sessions = state.sessions.list_sessions().await;
    Json(sessions.into_iter().map(api::SessionInfo::from).collect())
}

async fn get_session(
    State(state): State<AppState>,
    Path(id): Path<String>,
) -> Result<Json<api::SessionInfo>, (StatusCode, &'static str)> {
    let session_id = uuid::Uuid::parse_str(&id)
        .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID"))?;

    let session = state.sessions.get_session(session_id).await
        .ok_or((StatusCode::NOT_FOUND, "Session not found"))?;

    Ok(Json(api::SessionInfo::from(session)))
}

async fn disconnect_session(
    State(state): State<AppState>,
    Path(id): Path<String>,
) -> impl IntoResponse {
    let session_id = match uuid::Uuid::parse_str(&id) {
        Ok(id) => id,
        Err(_) => return (StatusCode::BAD_REQUEST, "Invalid session ID"),
    };

    if state.sessions.disconnect_session(session_id, "Disconnected by administrator").await {
        info!("Session {} disconnected by admin", session_id);
        (StatusCode::OK, "Session disconnected")
    } else {
        (StatusCode::NOT_FOUND, "Session not found")
    }
}

// Static page handlers
async fn serve_login() -> impl IntoResponse {
    match tokio::fs::read_to_string("static/login.html").await {
        Ok(content) => Html(content).into_response(),
        Err(_) => (StatusCode::NOT_FOUND, "Page not found").into_response(),
    }
}

async fn serve_dashboard() -> impl IntoResponse {
    match tokio::fs::read_to_string("static/dashboard.html").await {
        Ok(content) => Html(content).into_response(),
        Err(_) => (StatusCode::NOT_FOUND, "Page not found").into_response(),
    }
}


Found 1 file
guru-connect\agent\src\session\mod.rs


I need to:
1. Check the api module to see what SessionInfo looks like and update it
2. Add list_displays to the capture module

Let me check both.


Let me check the api module and capture module:


<tool_use_error>File does not exist. Did you mean api?</tool_use_error>

Error: File does not exist. Did you mean api?


     1→//! Screen capture module
     2→//!
     3→//! Provides DXGI Desktop Duplication for high-performance screen capture on Windows 8+,
     4→//! with GDI fallback for legacy systems or edge cases.
     5→
     6→#[cfg(windows)]
     7→mod dxgi;
     8→#[cfg(windows)]
     9→mod gdi;
    10→mod display;
    11→
    12→pub use display::{Display, DisplayInfo};
    13→
    14→use anyhow::Result;
    15→use std::time::Instant;
    16→
    17→/// Captured frame data
    18→#[derive(Debug)]
    19→pub struct CapturedFrame {
    20→    /// Frame width in pixels
    21→    pub width: u32,
    22→
    23→    /// Frame height in pixels
    24→    pub height: u32,
    25→
    26→    /// Raw BGRA pixel data (4 bytes per pixel)
    27→    pub data: Vec<u8>,
    28→
    29→    /// Timestamp when frame was captured
    30→    pub timestamp: Instant,
    31→
    32→    /// Display ID this frame is from
    33→    pub display_id: u32,
    34→
    35→    /// Regions that changed since last frame (if available)
    36→    pub dirty_rects: Option<Vec<DirtyRect>>,
    37→}
    38→
    39→/// Rectangular region that changed
    40→#[derive(Debug, Clone, Copy)]
    41→pub struct DirtyRect {
    42→    pub x: u32,
    43→    pub y: u32,
    44→    pub width: u32,
    45→    pub height: u32,
    46→}
    47→
    48→/// Screen capturer trait
    49→pub trait Capturer: Send {
    50→    /// Capture the next frame
    51→    ///
    52→    /// Returns None if no new frame is available (screen unchanged)
    53→    fn capture(&mut self) -> Result<Option<CapturedFrame>>;
    54→
    55→    /// Get the current display info
    56→    fn display(&self) -> &Display;
    57→
    58→    /// Check if capturer is still valid (display may have changed)
    59→    fn is_valid(&self) -> bool;
    60→}
    61→
    62→/// Create a capturer for the specified display
    63→#[cfg(windows)]
    64→pub fn create_capturer(display: Display, use_dxgi: bool, gdi_fallback: bool) -> Result<Box<dyn Capturer>> {
    65→    if use_dxgi {
    66→        match dxgi::DxgiCapturer::new(display.clone()) {
    67→            Ok(capturer) => {
    68→                tracing::info!("Using DXGI Desktop Duplication for capture");
    69→                return Ok(Box::new(capturer));
    70→            }
    71→            Err(e) => {
    72→                tracing::warn!("DXGI capture failed: {}, trying fallback", e);
    73→                if !gdi_fallback {
    74→                    return Err(e);
    75→                }
    76→            }
    77→        }
    78→    }
    79→
    80→    // GDI fallback
    81→    tracing::info!("Using GDI for capture");
    82→    Ok(Box::new(gdi::GdiCapturer::new(display)?))
    83→}
    84→
    85→#[cfg(not(windows))]
    86→pub fn create_capturer(_display: Display, _use_dxgi: bool, _gdi_fallback: bool) -> Result<Box<dyn Capturer>> {
    87→    anyhow::bail!("Screen capture only supported on Windows")
    88→}
    89→
    90→/// Get all available displays
    91→pub fn enumerate_displays() -> Result<Vec<Display>> {
    92→    display::enumerate_displays()
    93→}
    94→
    95→/// Get the primary display
    96→pub fn primary_display() -> Result<Display> {
    97→    let displays = enumerate_displays()?;
    98→    displays
    99→        .into_iter()
   100→        .find(|d| d.is_primary)
   101→        .ok_or_else(|| anyhow::anyhow!("No primary display found"))
   102→}
   103→

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

//! Screen capture module
//!
//! Provides DXGI Desktop Duplication for high-performance screen capture on Windows 8+,
//! with GDI fallback for legacy systems or edge cases.

#[cfg(windows)]
mod dxgi;
#[cfg(windows)]
mod gdi;
mod display;

pub use display::{Display, DisplayInfo};

use anyhow::Result;
use std::time::Instant;

/// Captured frame data
#[derive(Debug)]
pub struct CapturedFrame {
    /// Frame width in pixels
    pub width: u32,

    /// Frame height in pixels
    pub height: u32,

    /// Raw BGRA pixel data (4 bytes per pixel)
    pub data: Vec<u8>,

    /// Timestamp when frame was captured
    pub timestamp: Instant,

    /// Display ID this frame is from
    pub display_id: u32,

    /// Regions that changed since last frame (if available)
    pub dirty_rects: Option<Vec<DirtyRect>>,
}

/// Rectangular region that changed
#[derive(Debug, Clone, Copy)]
pub struct DirtyRect {
    pub x: u32,
    pub y: u32,
    pub width: u32,
    pub height: u32,
}

/// Screen capturer trait
pub trait Capturer: Send {
    /// Capture the next frame
    ///
    /// Returns None if no new frame is available (screen unchanged)
    fn capture(&mut self) -> Result<Option<CapturedFrame>>;

    /// Get the current display info
    fn display(&self) -> &Display;

    /// Check if capturer is still valid (display may have changed)
    fn is_valid(&self) -> bool;
}

/// Create a capturer for the specified display
#[cfg(windows)]
pub fn create_capturer(display: Display, use_dxgi: bool, gdi_fallback: bool) -> Result<Box<dyn Capturer>> {
    if use_dxgi {
        match dxgi::DxgiCapturer::new(display.clone()) {
            Ok(capturer) => {
                tracing::info!("Using DXGI Desktop Duplication for capture");
                return Ok(Box::new(capturer));
            }
            Err(e) => {
                tracing::warn!("DXGI capture failed: {}, trying fallback", e);
                if !gdi_fallback {
                    return Err(e);
                }
            }
        }
    }

    // GDI fallback
    tracing::info!("Using GDI for capture");
    Ok(Box::new(gdi::GdiCapturer::new(display)?))
}

#[cfg(not(windows))]
pub fn create_capturer(_display: Display, _use_dxgi: bool, _gdi_fallback: bool) -> Result<Box<dyn Capturer>> {
    anyhow::bail!("Screen capture only supported on Windows")
}

/// Get all available displays
pub fn enumerate_displays() -> Result<Vec<Display>> {
    display::enumerate_displays()
}

/// Get the primary display
pub fn primary_display() -> Result<Display> {
    let displays = enumerate_displays()?;
    displays
        .into_iter()
        .find(|d| d.is_primary)
        .ok_or_else(|| anyhow::anyhow!("No primary display found"))
}


I see that:1. The function is called `enumerate_displays()` not `list_displays()` - I need to fix the session module
2. There's no api.rs file - the SessionInfo is created inline. Let me check if there's an api module in the server directory


Let me fix the capture reference and check what api module exists:


Edit

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

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


No files found


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:
   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::enumerate_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 {

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

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

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

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


The server's main.rs references `api::SessionInfo` but there's no api.rs file. This suggests it was created previously. Let me search for where SessionInfo is defined and create it if needed.


Found 2 files
guru-connect\server\src\main.rs
guru-connect\server\src\api\mod.rs


The api module exists as api/mod.rs. Let me read it and update it.


     1→//! REST API endpoints
     2→
     3→use axum::{
     4→    extract::{Path, State},
     5→    Json,
     6→};
     7→use serde::Serialize;
     8→use uuid::Uuid;
     9→
    10→use crate::session::SessionManager;
    11→
    12→/// Session info returned by API
    13→#[derive(Debug, Serialize)]
    14→pub struct SessionInfo {
    15→    pub id: String,
    16→    pub agent_id: String,
    17→    pub agent_name: String,
    18→    pub started_at: String,
    19→    pub viewer_count: usize,
    20→}
    21→
    22→impl From<crate::session::Session> for SessionInfo {
    23→    fn from(s: crate::session::Session) -> Self {
    24→        Self {
    25→            id: s.id.to_string(),
    26→            agent_id: s.agent_id,
    27→            agent_name: s.agent_name,
    28→            started_at: s.started_at.to_rfc3339(),
    29→            viewer_count: s.viewer_count,
    30→        }
    31→    }
    32→}
    33→
    34→/// List all active sessions
    35→pub async fn list_sessions(
    36→    State(sessions): State<SessionManager>,
    37→) -> Json<Vec<SessionInfo>> {
    38→    let sessions = sessions.list_sessions().await;
    39→    Json(sessions.into_iter().map(SessionInfo::from).collect())
    40→}
    41→
    42→/// Get a specific session by ID
    43→pub async fn get_session(
    44→    State(sessions): State<SessionManager>,
    45→    Path(id): Path<String>,
    46→) -> Result<Json<SessionInfo>, (axum::http::StatusCode, &'static str)> {
    47→    let session_id = Uuid::parse_str(&id)
    48→        .map_err(|_| (axum::http::StatusCode::BAD_REQUEST, "Invalid session ID"))?;
    49→
    50→    let session = sessions.get_session(session_id).await
    51→        .ok_or((axum::http::StatusCode::NOT_FOUND, "Session not found"))?;
    52→
    53→    Ok(Json(SessionInfo::from(session)))
    54→}
    55→

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

//! REST API endpoints

use axum::{
    extract::{Path, State},
    Json,
};
use serde::Serialize;
use uuid::Uuid;

use crate::session::SessionManager;

/// Session info returned by API
#[derive(Debug, Serialize)]
pub struct SessionInfo {
    pub id: String,
    pub agent_id: String,
    pub agent_name: String,
    pub started_at: String,
    pub viewer_count: usize,
}

impl From<crate::session::Session> for SessionInfo {
    fn from(s: crate::session::Session) -> Self {
        Self {
            id: s.id.to_string(),
            agent_id: s.agent_id,
            agent_name: s.agent_name,
            started_at: s.started_at.to_rfc3339(),
            viewer_count: s.viewer_count,
        }
    }
}

/// List all active sessions
pub async fn list_sessions(
    State(sessions): State<SessionManager>,
) -> Json<Vec<SessionInfo>> {
    let sessions = sessions.list_sessions().await;
    Json(sessions.into_iter().map(SessionInfo::from).collect())
}

/// Get a specific session by ID
pub async fn get_session(
    State(sessions): State<SessionManager>,
    Path(id): Path<String>,
) -> Result<Json<SessionInfo>, (axum::http::StatusCode, &'static str)> {
    let session_id = Uuid::parse_str(&id)
        .map_err(|_| (axum::http::StatusCode::BAD_REQUEST, "Invalid session ID"))?;

    let session = sessions.get_session(session_id).await
        .ok_or((axum::http::StatusCode::NOT_FOUND, "Session not found"))?;

    Ok(Json(SessionInfo::from(session)))
}


I need to update SessionInfo to include the new fields from Session (is_streaming, last_heartbeat, etc.).


Let me update SessionInfo with the new fields: