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

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

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


I need to add a way to disconnect persistent agents. Looking at the code:

1. The session manager stores sessions with frame_tx and input_tx channels
2. To disconnect an agent, I need to send a Disconnect message to them via the input_tx channel
3. Then the agent's handler will receive it, process it, and disconnect

The problem is that the current structure doesn't expose a way to send arbitrary messages to agents. The input_tx is meant for input events from viewers.

I have two options:
1. Add a separate channel for control messages to agents
2. Reuse the input_tx to send disconnect messages (the agent should handle them)

Actually, looking more closely, the input_tx sends data to the agent via the WebSocket. If I encode a Disconnect message and send it through input_tx, it should reach the agent.

Let me add:
1. A method in SessionManager to disconnect a session (sends Disconnect message)
2. An API endpoint DELETE /api/sessions/:id
3. Update the dashboard UI to add a Disconnect button

Let me start with the SessionManager.


I need to add a way to send disconnect messages to agents. Let me update the session manager first:


Edit

    /// Remove a session (when agent disconnects)
    pub async fn remove_session(&self, session_id: SessionId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.remove(&session_id) {
            let mut agents = self.agents.write().await;
            agents.remove(&session_data.info.agent_id);
        }
    }

    /// Remove a session (when agent disconnects)
    pub async fn remove_session(&self, session_id: SessionId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.remove(&session_id) {
            let mut agents = self.agents.write().await;
            agents.remove(&session_data.info.agent_id);
        }
    }

    /// Disconnect a session by sending a disconnect message to the agent
    /// Returns true if the message was sent successfully
    pub async fn disconnect_session(&self, session_id: SessionId, reason: &str) -> bool {
        let sessions = self.sessions.read().await;
        if let Some(session_data) = sessions.get(&session_id) {
            // Create disconnect message
            use crate::proto;
            use prost::Message;

            let disconnect_msg = proto::Message {
                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
                    reason: reason.to_string(),
                })),
            };

            let mut buf = Vec::new();
            if disconnect_msg.encode(&mut buf).is_ok() {
                // Send via input channel (will be forwarded to agent's WebSocket)
                if session_data.input_tx.send(buf).await.is_ok() {
                    return true;
                }
            }
        }
        false
    }


The file C:\Users\MikeSwanson\Claude\guru-connect\server\src\session\mod.rs has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
   124→            session_data.info.viewer_count = session_data.info.viewer_count.saturating_sub(1);
   125→        }
   126→    }
   127→
   128→    /// Remove a session (when agent disconnects)
   129→    pub async fn remove_session(&self, session_id: SessionId) {
   130→        let mut sessions = self.sessions.write().await;
   131→        if let Some(session_data) = sessions.remove(&session_id) {
   132→            let mut agents = self.agents.write().await;
   133→            agents.remove(&session_data.info.agent_id);
   134→        }
   135→    }
   136→
   137→    /// Disconnect a session by sending a disconnect message to the agent
   138→    /// Returns true if the message was sent successfully
   139→    pub async fn disconnect_session(&self, session_id: SessionId, reason: &str) -> bool {
   140→        let sessions = self.sessions.read().await;
   141→        if let Some(session_data) = sessions.get(&session_id) {
   142→            // Create disconnect message
   143→            use crate::proto;
   144→            use prost::Message;
   145→
   146→            let disconnect_msg = proto::Message {
   147→                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
   148→                    reason: reason.to_string(),
   149→                })),
   150→            };
   151→
   152→            let mut buf = Vec::new();
   153→            if disconnect_msg.encode(&mut buf).is_ok() {
   154→                // Send via input channel (will be forwarded to agent's WebSocket)
   155→                if session_data.input_tx.send(buf).await.is_ok() {
   156→                    return true;
   157→                }
   158→            }
   159→        }
   160→        false
   161→    }
   162→
   163→    /// List all active sessions
   164→    pub async fn list_sessions(&self) -> Vec<Session> {
   165→        let sessions = self.sessions.read().await;

    /// Remove a session (when agent disconnects)
    pub async fn remove_session(&self, session_id: SessionId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.remove(&session_id) {
            let mut agents = self.agents.write().await;
            agents.remove(&session_data.info.agent_id);
        }
    }

    /// Disconnect a session by sending a disconnect message to the agent
    /// Returns true if the message was sent successfully
    pub async fn disconnect_session(&self, session_id: SessionId, reason: &str) -> bool {
        let sessions = self.sessions.read().await;
        if let Some(session_data) = sessions.get(&session_id) {
            // Create disconnect message
            use crate::proto;
            use prost::Message;

            let disconnect_msg = proto::Message {
                payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
                    reason: reason.to_string(),
                })),
            };

            let mut buf = Vec::new();
            if disconnect_msg.encode(&mut buf).is_ok() {
                // Send via input channel (will be forwarded to agent's WebSocket)
                if session_data.input_tx.send(buf).await.is_ok() {
                    return true;
                }
            }
        }
        false
    }

//! Session management for GuruConnect
//!
//! Manages active remote desktop sessions, tracking which agents
//! are connected and which viewers are watching them.

use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{broadcast, RwLock};
use uuid::Uuid;

/// Unique identifier for a session
pub type SessionId = Uuid;

/// Unique identifier for an agent
pub type AgentId = String;

/// Session state
#[derive(Debug, Clone)]
pub struct Session {
    pub id: SessionId,
    pub agent_id: AgentId,
    pub agent_name: String,
    pub started_at: chrono::DateTime<chrono::Utc>,
    pub viewer_count: usize,
}

/// Channel for sending frames from agent to viewers
pub type FrameSender = broadcast::Sender<Vec<u8>>;
pub type FrameReceiver = broadcast::Receiver<Vec<u8>>;

/// Channel for sending input events from viewer to agent
pub type InputSender = tokio::sync::mpsc::Sender<Vec<u8>>;
pub type InputReceiver = tokio::sync::mpsc::Receiver<Vec<u8>>;

/// Internal session data with channels
struct SessionData {
    info: Session,
    /// Channel for video frames (agent -> viewers)
    frame_tx: FrameSender,
    /// Channel for input events (viewer -> agent)
    input_tx: InputSender,
    input_rx: Option<InputReceiver>,
}

/// Manages all active sessions
#[derive(Clone)]
pub struct SessionManager {
    sessions: Arc<RwLock<HashMap<SessionId, SessionData>>>,
    agents: Arc<RwLock<HashMap<AgentId, SessionId>>>,
}

impl SessionManager {
    pub fn new() -> Self {
        Self {
            sessions: Arc::new(RwLock::new(HashMap::new())),
            agents: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Register a new agent and create a session
    pub async fn register_agent(&self, agent_id: AgentId, agent_name: String) -> (SessionId, FrameSender, InputReceiver) {
        let session_id = Uuid::new_v4();

        // Create channels
        let (frame_tx, _) = broadcast::channel(16); // Buffer 16 frames
        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64); // Buffer 64 input events

        let session = Session {
            id: session_id,
            agent_id: agent_id.clone(),
            agent_name,
            started_at: chrono::Utc::now(),
            viewer_count: 0,
        };

        let session_data = SessionData {
            info: session,
            frame_tx: frame_tx.clone(),
            input_tx,
            input_rx: None, // Will be taken by the agent handler
        };

        let mut sessions = self.sessions.write().await;
        sessions.insert(session_id, session_data);

        let mut agents = self.agents.write().await;
        agents.insert(agent_id, session_id);

        (session_id, frame_tx, input_rx)
    }

    /// Get a session by agent ID
    pub async fn get_session_by_agent(&self, agent_id: &str) -> Option<Session> {
        let agents = self.agents.read().await;
        let session_id = agents.get(agent_id)?;

        let sessions = self.sessions.read().await;
        sessions.get(session_id).map(|s| s.info.clone())
    }

    /// Get a session by session ID
    pub async fn get_session(&self, session_id: SessionId) -> Option<Session> {
        let sessions = self.sessions.read().await;
        sessions.get(&session_id).map(|s| s.info.clone())
    }

    /// Join a session as a viewer
    pub async fn join_session(&self, session_id: SessionId) -> Option<(FrameReceiver, InputSender)> {
        let mut sessions = self.sessions.write().await;
        let session_data = sessions.get_mut(&session_id)?;

        session_data.info.viewer_count += 1;

        let frame_rx = session_data.frame_tx.subscribe();
        let input_tx = session_data.input_tx.clone();

        Some((frame_rx, input_tx))
    }

    /// Leave a session as a viewer
    pub async fn leave_session(&self, session_id: SessionId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.get_mut(&session_id) {
            session_data.info.viewer_count = session_data.info.viewer_count.saturating_sub(1);
        }
    }

    /// Remove a session (when agent disconnects)
    pub async fn remove_session(&self, session_id: SessionId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.remove(&session_id) {
            let mut agents = self.agents.write().await;
            agents.remove(&session_data.info.agent_id);
        }
    }

    /// List all active sessions
    pub async fn list_sessions(&self) -> Vec<Session> {
        let sessions = self.sessions.read().await;
        sessions.values().map(|s| s.info.clone()).collect()
    }
}

impl Default for SessionManager {
    fn default() -> Self {
        Self::new()
    }
}

+    /// Disconnect a session by sending a disconnect message to the agent

+    /// Returns true if the message was sent successfully

+    pub async fn disconnect_session(&self, session_id: SessionId, reason: &str) -> bool {

+        let sessions = self.sessions.read().await;

+        if let Some(session_data) = sessions.get(&session_id) {

+            // Create disconnect message

+            let disconnect_msg = proto::Message {

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

+                    reason: reason.to_string(),

+            if disconnect_msg.encode(&mut buf).is_ok() {

+                // Send via input channel (will be forwarded to agent's WebSocket)

+                if session_data.input_tx.send(buf).await.is_ok() {


Now I need to add an API endpoint in main.rs to call this disconnect method. Let me add a DELETE endpoint for sessions.


Now let me add the API endpoint in main.rs:


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

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


The file C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
    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))

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

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

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

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

+        .route("/api/sessions/:id", axum::routing::delete(disconnect_session))


Now I need to add the disconnect_session handler function. Let me add it near the other session handlers.


Now let me add the handler function:


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 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")
    }
}


The file C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
   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 {

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")
    }
}

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

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

+    let session_id = match uuid::Uuid::parse_str(&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")

+        (StatusCode::NOT_FOUND, "Session not found")


Now I need to update the dashboard.html to add a Disconnect button in the Access tab. Let me find the relevant section in the dashboard.


Now let me update the dashboard to add a Disconnect button in the Access tab:


'
' +
                    '
Actions
' +
                    '<button class="btn btn-primary" style="width: 100%; margin-bottom: 8px;" onclick="connectToMachine(\'' + m.id + '\')">Connect</button>' +
                    '<button class="btn btn-outline" style="width: 100%;" disabled>Transfer Files</button>' +
                '
';

'
' +
                    '
Actions
' +
                    '<button class="btn btn-primary" style="width: 100%; margin-bottom: 8px;" onclick="connectToMachine(\'' + m.id + '\')">Connect</button>' +
                    '<button class="btn btn-outline" style="width: 100%; margin-bottom: 8px;" onclick="openChat(\'' + m.id + '\', \'' + (m.agent_name || 'Client').replace(/'/g, "\\'") + '\')">Chat</button>' +
                    '<button class="btn btn-outline" style="width: 100%; margin-bottom: 8px;" disabled>Transfer Files</button>' +
                    '<button class="btn btn-outline" style="width: 100%; color: hsl(0, 62.8%, 50%);" onclick="disconnectMachine(\'' + m.id + '\', \'' + (m.agent_name || m.agent_id).replace(/'/g, "\\'") + '\')">Disconnect</button>' +
                '
';


The file C:\Users\MikeSwanson\Claude\guru-connect\server\static\dashboard.html has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
   705→                    '
Session ID' + m.id.slice(0,8) + '...
' +
   706→                    '
Connected' + started + '
' +
   707→                    '
Viewers' + m.viewer_count + '
' +
   708→                '
' +
   709→                '
' +
   710→                    '
Actions
' +
   711→                    '<button class="btn btn-primary" style="width: 100%; margin-bottom: 8px;" onclick="connectToMachine(\'' + m.id + '\')">Connect</button>' +
   712→                    '<button class="btn btn-outline" style="width: 100%; margin-bottom: 8px;" onclick="openChat(\'' + m.id + '\', \'' + (m.agent_name || 'Client').replace(/'/g, "\\'") + '\')">Chat</button>' +
   713→                    '<button class="btn btn-outline" style="width: 100%; margin-bottom: 8px;" disabled>Transfer Files</button>' +
   714→                    '<button class="btn btn-outline" style="width: 100%; color: hsl(0, 62.8%, 50%);" onclick="disconnectMachine(\'' + m.id + '\', \'' + (m.agent_name || m.agent_id).replace(/'/g, "\\'") + '\')">Disconnect</button>' +
   715→                '
';
   716→        }
   717→
   718→        function connectToMachine(sessionId) {
   719→            // TODO: Open viewer in new window

'
' +
                    '
Actions
' +
                    '<button class="btn btn-primary" style="width: 100%; margin-bottom: 8px;" onclick="connectToMachine(\'' + m.id + '\')">Connect</button>' +
                    '<button class="btn btn-outline" style="width: 100%;" disabled>Transfer Files</button>' +
                '
';

'
' +
                    '
Actions
' +
                    '<button class="btn btn-primary" style="width: 100%; margin-bottom: 8px;" onclick="connectToMachine(\'' + m.id + '\')">Connect</button>' +
                    '<button class="btn btn-outline" style="width: 100%; margin-bottom: 8px;" onclick="openChat(\'' + m.id + '\', \'' + (m.agent_name || 'Client').replace(/'/g, "\\'") + '\')">Chat</button>' +
                    '<button class="btn btn-outline" style="width: 100%; margin-bottom: 8px;" disabled>Transfer Files</button>' +
                    '<button class="btn btn-outline" style="width: 100%; color: hsl(0, 62.8%, 50%);" onclick="disconnectMachine(\'' + m.id + '\', \'' + (m.agent_name || m.agent_id).replace(/'/g, "\\'") + '\')">Disconnect</button>' +
                '
';

<!DOCTYPE html>

    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>GuruConnect - Dashboard</title>

GuruConnect

            Loading...
            <button class="logout-btn" id="logoutBtn">Sign Out</button>

        <button class="tab active" data-tab="support">Support</button>
        <button class="tab" data-tab="access">Access</button>
        <button class="tab" data-tab="build">Build</button>
        <button class="tab" data-tab="settings">Settings</button>

Active Support Sessions

Temporary sessions initiated by support codes

                    <button class="btn btn-primary" id="generateCodeBtn">Generate Code</button>

                                Code
                                Status
                                Created
                                Technician
                                Actions

No active sessions

Generate a code to start a support session

Status

                            All Machines
                            0

                            Online
                            0

                            Offline
                            0

By Company

No machines installed

No machines

Install the agent on a machine to see it here

Select a machine

Click a machine to view details

Installer Builder

Create customized agent installers for unattended access

                        Name
                        <input type="text" id="buildName" placeholder="Machine name (auto if blank)">

                        Company
                        <input type="text" id="buildCompany" placeholder="Client organization">

                        Site
                        <input type="text" id="buildSite" placeholder="Physical location">

                        Department
                        <input type="text" id="buildDepartment" placeholder="Business unit">

                        Device Type
                        <select id="buildDeviceType">
                            <option value="workstation">Workstation</option>
                            <option value="laptop">Laptop</option>
                            <option value="server">Server</option>
                        </select>

                        Tag
                        <input type="text" id="buildTag" placeholder="Custom label">

                    <button class="btn btn-primary" disabled>Build EXE (64-bit)</button>
                    <button class="btn btn-outline" disabled>Build EXE (32-bit)</button>
                    <button class="btn btn-outline" disabled>Build MSI</button>

                    Agent builds will be available once the agent is compiled.

Settings

Configure your GuruConnect preferences

                        Theme
                        <select disabled>
                            <option>Dark (Default)</option>
                            <option>Light</option>
                            <option>System</option>
                        </select>

                        Notifications
                        <select disabled>
                            <option>All notifications</option>
                            <option>Important only</option>
                            <option>None</option>
                        </select>

                    Additional settings coming soon.

Chat with Client

                <button class="modal-close" id="chatClose">×</button>

No messages yet. Start the conversation!

                <input type="text" class="chat-input" id="chatInput" placeholder="Type a message..." />
                <button class="chat-send" id="chatSend">Send</button>

-                    '<button class="btn btn-outline" style="width: 100%;" disabled>Transfer Files</button>' +


And be sure to trigger a full uninstall on the client side for ended sessions.


Now I need to add the disconnectMachine JavaScript function. Let me add it after the connectToMachine function.


Now let me add the disconnectMachine JavaScript function:


        function connectToMachine(sessionId) {
            // TODO: Open viewer in new window
            alert("Viewer not yet implemented.\\n\\nSession ID: " + sessionId + "\\n\\nWebSocket: wss://connect.azcomputerguru.com/ws/viewer?session_id=" + sessionId);
        }

        // Refresh machines every 5 seconds

        function connectToMachine(sessionId) {
            // TODO: Open viewer in new window
            alert("Viewer not yet implemented.\\n\\nSession ID: " + sessionId + "\\n\\nWebSocket: wss://connect.azcomputerguru.com/ws/viewer?session_id=" + sessionId);
        }

        async function disconnectMachine(sessionId, machineName) {
            if (!confirm("Disconnect " + machineName + "?\\n\\nThis will end the remote session.")) return;
            try {
                const response = await fetch("/api/sessions/" + sessionId, { method: "DELETE" });
                if (response.ok) {
                    selectedMachine = null;
                    renderMachineDetail();
                    loadMachines();
                } else {
                    alert("Failed to disconnect: " + await response.text());
                }
            } catch (err) {
                alert("Error disconnecting machine");
            }
        }

        // Refresh machines every 5 seconds


The file C:\Users\MikeSwanson\Claude\guru-connect\server\static\dashboard.html has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
   714→                    '<button class="btn btn-outline" style="width: 100%; color: hsl(0, 62.8%, 50%);" onclick="disconnectMachine(\'' + m.id + '\', \'' + (m.agent_name || m.agent_id).replace(/'/g, "\\'") + '\')">Disconnect</button>' +
   715→                '</div>';
   716→        }
   717→
   718→        function connectToMachine(sessionId) {
   719→            // TODO: Open viewer in new window
   720→            alert("Viewer not yet implemented.\\n\\nSession ID: " + sessionId + "\\n\\nWebSocket: wss://connect.azcomputerguru.com/ws/viewer?session_id=" + sessionId);
   721→        }
   722→
   723→        async function disconnectMachine(sessionId, machineName) {
   724→            if (!confirm("Disconnect " + machineName + "?\\n\\nThis will end the remote session.")) return;
   725→            try {
   726→                const response = await fetch("/api/sessions/" + sessionId, { method: "DELETE" });
   727→                if (response.ok) {
   728→                    selectedMachine = null;
   729→                    renderMachineDetail();
   730→                    loadMachines();
   731→                } else {
   732→                    alert("Failed to disconnect: " + await response.text());
   733→                }
   734→            } catch (err) {
   735→                alert("Error disconnecting machine");
   736→            }
   737→        }
   738→
   739→        // Refresh machines every 5 seconds
   740→        loadMachines();
   741→        setInterval(loadMachines, 5000);
   742→
   743→        // ========== Chat Functions ==========

        function connectToMachine(sessionId) {
            // TODO: Open viewer in new window
            alert("Viewer not yet implemented.\\n\\nSession ID: " + sessionId + "\\n\\nWebSocket: wss://connect.azcomputerguru.com/ws/viewer?session_id=" + sessionId);
        }

        // Refresh machines every 5 seconds

        function connectToMachine(sessionId) {
            // TODO: Open viewer in new window
            alert("Viewer not yet implemented.\\n\\nSession ID: " + sessionId + "\\n\\nWebSocket: wss://connect.azcomputerguru.com/ws/viewer?session_id=" + sessionId);
        }

        async function disconnectMachine(sessionId, machineName) {
            if (!confirm("Disconnect " + machineName + "?\\n\\nThis will end the remote session.")) return;
            try {
                const response = await fetch("/api/sessions/" + sessionId, { method: "DELETE" });
                if (response.ok) {
                    selectedMachine = null;
                    renderMachineDetail();
                    loadMachines();
                } else {
                    alert("Failed to disconnect: " + await response.text());
                }
            } catch (err) {
                alert("Error disconnecting machine");
            }
        }

        // Refresh machines every 5 seconds

+        async function disconnectMachine(sessionId, machineName) {

+            if (!confirm("Disconnect " + machineName + "?\\n\\nThis will end the remote session.")) return;

+                const response = await fetch("/api/sessions/" + sessionId, { method: "DELETE" });

+                    selectedMachine = null;

+                    renderMachineDetail();

+                    alert("Failed to disconnect: " + await response.text());

+                alert("Error disconnecting machine");


The user wants the agent to uninstall itself when a persistent session is ended (similar to how support sessions clean up). Let me check how the agent currently handles disconnect messages, then add uninstall logic.Currently, for support sessions, the agent:
1. Removes itself from startup registry2. Shows a message box
3. Exits cleanlyFor persistent sessions being disconnected by admin, I need to add similar cleanup. Let me check the agent's main.rs to see how disconnections are handled.


I need to update the agent to handle disconnect messages and clean itself up. Let me check and update the agent's disconnect handling:


     1→//! Startup persistence for the agent
     2→//!
     3→//! Handles adding/removing the agent from Windows startup.
     4→
     5→use anyhow::Result;
     6→use tracing::{info, warn, error};
     7→
     8→#[cfg(windows)]
     9→use windows::Win32::System::Registry::{
    10→    RegOpenKeyExW, RegSetValueExW, RegDeleteValueW, RegCloseKey,
    11→    HKEY_CURRENT_USER, KEY_WRITE, REG_SZ,
    12→};
    13→#[cfg(windows)]
    14→use windows::core::PCWSTR;
    15→
    16→const STARTUP_KEY: &str = r"Software\Microsoft\Windows\CurrentVersion\Run";
    17→const STARTUP_VALUE_NAME: &str = "GuruConnect";
    18→
    19→/// Add the current executable to Windows startup
    20→#[cfg(windows)]
    21→pub fn add_to_startup() -> Result<()> {
    22→    use std::ffi::OsStr;
    23→    use std::os::windows::ffi::OsStrExt;
    24→
    25→    // Get the path to the current executable
    26→    let exe_path = std::env::current_exe()?;
    27→    let exe_path_str = exe_path.to_string_lossy();
    28→
    29→    info!("Adding to startup: {}", exe_path_str);
    30→
    31→    // Convert strings to wide strings
    32→    let key_path: Vec<u16> = OsStr::new(STARTUP_KEY)
    33→        .encode_wide()
    34→        .chain(std::iter::once(0))
    35→        .collect();
    36→    let value_name: Vec<u16> = OsStr::new(STARTUP_VALUE_NAME)
    37→        .encode_wide()
    38→        .chain(std::iter::once(0))
    39→        .collect();
    40→    let value_data: Vec<u16> = OsStr::new(&*exe_path_str)
    41→        .encode_wide()
    42→        .chain(std::iter::once(0))
    43→        .collect();
    44→
    45→    unsafe {
    46→        let mut hkey = windows::Win32::Foundation::HANDLE::default();
    47→
    48→        // Open the Run key
    49→        let result = RegOpenKeyExW(
    50→            HKEY_CURRENT_USER,
    51→            PCWSTR(key_path.as_ptr()),
    52→            0,
    53→            KEY_WRITE,
    54→            &mut hkey as *mut _ as *mut _,
    55→        );
    56→
    57→        if result.is_err() {
    58→            anyhow::bail!("Failed to open registry key: {:?}", result);
    59→        }
    60→
    61→        let hkey_raw = std::mem::transmute::<_, windows::Win32::System::Registry::HKEY>(hkey);
    62→
    63→        // Set the value
    64→        let data_bytes = std::slice::from_raw_parts(
    65→            value_data.as_ptr() as *const u8,
    66→            value_data.len() * 2,
    67→        );
    68→
    69→        let set_result = RegSetValueExW(
    70→            hkey_raw,
    71→            PCWSTR(value_name.as_ptr()),
    72→            0,
    73→            REG_SZ,
    74→            Some(data_bytes),
    75→        );
    76→
    77→        let _ = RegCloseKey(hkey_raw);
    78→
    79→        if set_result.is_err() {
    80→            anyhow::bail!("Failed to set registry value: {:?}", set_result);
    81→        }
    82→    }
    83→
    84→    info!("Successfully added to startup");
    85→    Ok(())
    86→}
    87→
    88→/// Remove the agent from Windows startup
    89→#[cfg(windows)]
    90→pub fn remove_from_startup() -> Result<()> {
    91→    use std::ffi::OsStr;
    92→    use std::os::windows::ffi::OsStrExt;
    93→
    94→    info!("Removing from startup");
    95→
    96→    let key_path: Vec<u16> = OsStr::new(STARTUP_KEY)
    97→        .encode_wide()
    98→        .chain(std::iter::once(0))
    99→        .collect();
   100→    let value_name: Vec<u16> = OsStr::new(STARTUP_VALUE_NAME)
   101→        .encode_wide()
   102→        .chain(std::iter::once(0))
   103→        .collect();
   104→
   105→    unsafe {
   106→        let mut hkey = windows::Win32::Foundation::HANDLE::default();
   107→
   108→        let result = RegOpenKeyExW(
   109→            HKEY_CURRENT_USER,
   110→            PCWSTR(key_path.as_ptr()),
   111→            0,
   112→            KEY_WRITE,
   113→            &mut hkey as *mut _ as *mut _,
   114→        );
   115→
   116→        if result.is_err() {
   117→            warn!("Failed to open registry key for removal: {:?}", result);
   118→            return Ok(()); // Not an error if key doesn't exist
   119→        }
   120→
   121→        let hkey_raw = std::mem::transmute::<_, windows::Win32::System::Registry::HKEY>(hkey);
   122→
   123→        let delete_result = RegDeleteValueW(hkey_raw, PCWSTR(value_name.as_ptr()));
   124→
   125→        let _ = RegCloseKey(hkey_raw);
   126→
   127→        if delete_result.is_err() {
   128→            warn!("Registry value may not exist: {:?}", delete_result);
   129→        } else {
   130→            info!("Successfully removed from startup");
   131→        }
   132→    }
   133→
   134→    Ok(())
   135→}
   136→
   137→#[cfg(not(windows))]
   138→pub fn add_to_startup() -> Result<()> {
   139→    warn!("Startup persistence not implemented for this platform");
   140→    Ok(())
   141→}
   142→
   143→#[cfg(not(windows))]
   144→pub fn remove_from_startup() -> Result<()> {
   145→    Ok(())
   146→}
   147→

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

//! Startup persistence for the agent
//!
//! Handles adding/removing the agent from Windows startup.

use anyhow::Result;
use tracing::{info, warn, error};

#[cfg(windows)]
use windows::Win32::System::Registry::{
    RegOpenKeyExW, RegSetValueExW, RegDeleteValueW, RegCloseKey,
    HKEY_CURRENT_USER, KEY_WRITE, REG_SZ,
};
#[cfg(windows)]
use windows::core::PCWSTR;

const STARTUP_KEY: &str = r"Software\Microsoft\Windows\CurrentVersion\Run";
const STARTUP_VALUE_NAME: &str = "GuruConnect";

/// Add the current executable to Windows startup
#[cfg(windows)]
pub fn add_to_startup() -> Result<()> {
    use std::ffi::OsStr;
    use std::os::windows::ffi::OsStrExt;

    // Get the path to the current executable
    let exe_path = std::env::current_exe()?;
    let exe_path_str = exe_path.to_string_lossy();

    info!("Adding to startup: {}", exe_path_str);

    // Convert strings to wide strings
    let key_path: Vec<u16> = OsStr::new(STARTUP_KEY)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();
    let value_name: Vec<u16> = OsStr::new(STARTUP_VALUE_NAME)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();
    let value_data: Vec<u16> = OsStr::new(&*exe_path_str)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    unsafe {
        let mut hkey = windows::Win32::Foundation::HANDLE::default();

        // Open the Run key
        let result = RegOpenKeyExW(
            HKEY_CURRENT_USER,
            PCWSTR(key_path.as_ptr()),
            0,
            KEY_WRITE,
            &mut hkey as *mut _ as *mut _,
        );

        if result.is_err() {
            anyhow::bail!("Failed to open registry key: {:?}", result);
        }

        let hkey_raw = std::mem::transmute::<_, windows::Win32::System::Registry::HKEY>(hkey);

        // Set the value
        let data_bytes = std::slice::from_raw_parts(
            value_data.as_ptr() as *const u8,
            value_data.len() * 2,
        );

        let set_result = RegSetValueExW(
            hkey_raw,
            PCWSTR(value_name.as_ptr()),
            0,
            REG_SZ,
            Some(data_bytes),
        );

        let _ = RegCloseKey(hkey_raw);

        if set_result.is_err() {
            anyhow::bail!("Failed to set registry value: {:?}", set_result);
        }
    }

    info!("Successfully added to startup");
    Ok(())
}

/// Remove the agent from Windows startup
#[cfg(windows)]
pub fn remove_from_startup() -> Result<()> {
    use std::ffi::OsStr;
    use std::os::windows::ffi::OsStrExt;

    info!("Removing from startup");

    let key_path: Vec<u16> = OsStr::new(STARTUP_KEY)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();
    let value_name: Vec<u16> = OsStr::new(STARTUP_VALUE_NAME)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    unsafe {
        let mut hkey = windows::Win32::Foundation::HANDLE::default();

        let result = RegOpenKeyExW(
            HKEY_CURRENT_USER,
            PCWSTR(key_path.as_ptr()),
            0,
            KEY_WRITE,
            &mut hkey as *mut _ as *mut _,
        );

        if result.is_err() {
            warn!("Failed to open registry key for removal: {:?}", result);
            return Ok(()); // Not an error if key doesn't exist
        }

        let hkey_raw = std::mem::transmute::<_, windows::Win32::System::Registry::HKEY>(hkey);

        let delete_result = RegDeleteValueW(hkey_raw, PCWSTR(value_name.as_ptr()));

        let _ = RegCloseKey(hkey_raw);

        if delete_result.is_err() {
            warn!("Registry value may not exist: {:?}", delete_result);
        } else {
            info!("Successfully removed from startup");
        }
    }

    Ok(())
}

#[cfg(not(windows))]
pub fn add_to_startup() -> Result<()> {
    warn!("Startup persistence not implemented for this platform");
    Ok(())
}

#[cfg(not(windows))]
pub fn remove_from_startup() -> Result<()> {
    Ok(())
}


I need to add an uninstall function to the startup.rs module that will:
1. Remove from startup registry
2. Delete the executable (schedule for deletion on reboot if currently running)Then update the agent to call this when it receives a disconnect message from an administrator.Let me add the uninstall function to startup.rs and update the session handling.


Let me add an uninstall function to the startup module:


#[cfg(not(windows))]
pub fn add_to_startup() -> Result<()> {
    warn!("Startup persistence not implemented for this platform");
    Ok(())
}

#[cfg(not(windows))]
pub fn remove_from_startup() -> Result<()> {
    Ok(())
}

/// Full uninstall: remove from startup and delete the executable
#[cfg(windows)]
pub fn uninstall() -> Result<()> {
    use std::ffi::OsStr;
    use std::os::windows::ffi::OsStrExt;
    use windows::Win32::Storage::FileSystem::{MoveFileExW, MOVEFILE_DELAY_UNTIL_REBOOT};

    info!("Uninstalling agent");

    // First remove from startup
    let _ = remove_from_startup();

    // Get the path to the current executable
    let exe_path = std::env::current_exe()?;
    let exe_path_str = exe_path.to_string_lossy();

    info!("Scheduling deletion of: {}", exe_path_str);

    // Convert path to wide string
    let exe_wide: Vec<u16> = OsStr::new(&*exe_path_str)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    // Schedule the file for deletion on next reboot
    // This is necessary because the executable is currently running
    unsafe {
        let result = MoveFileExW(
            PCWSTR(exe_wide.as_ptr()),
            PCWSTR::null(),
            MOVEFILE_DELAY_UNTIL_REBOOT,
        );

        if result.is_err() {
            warn!("Failed to schedule file deletion: {:?}. File may need manual removal.", result);
        } else {
            info!("Executable scheduled for deletion on reboot");
        }
    }

    Ok(())
}

#[cfg(not(windows))]
pub fn add_to_startup() -> Result<()> {
    warn!("Startup persistence not implemented for this platform");
    Ok(())
}

#[cfg(not(windows))]
pub fn remove_from_startup() -> Result<()> {
    Ok(())
}

#[cfg(not(windows))]
pub fn uninstall() -> Result<()> {
    warn!("Uninstall not implemented for this platform");
    Ok(())
}


The file C:\Users\MikeSwanson\Claude\guru-connect\agent\src\startup.rs has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
   133→
   134→    Ok(())
   135→}
   136→
   137→/// Full uninstall: remove from startup and delete the executable
   138→#[cfg(windows)]
   139→pub fn uninstall() -> Result<()> {
   140→    use std::ffi::OsStr;
   141→    use std::os::windows::ffi::OsStrExt;
   142→    use windows::Win32::Storage::FileSystem::{MoveFileExW, MOVEFILE_DELAY_UNTIL_REBOOT};
   143→
   144→    info!("Uninstalling agent");
   145→
   146→    // First remove from startup
   147→    let _ = remove_from_startup();
   148→
   149→    // Get the path to the current executable
   150→    let exe_path = std::env::current_exe()?;
   151→    let exe_path_str = exe_path.to_string_lossy();
   152→
   153→    info!("Scheduling deletion of: {}", exe_path_str);
   154→
   155→    // Convert path to wide string
   156→    let exe_wide: Vec<u16> = OsStr::new(&*exe_path_str)
   157→        .encode_wide()
   158→        .chain(std::iter::once(0))
   159→        .collect();
   160→
   161→    // Schedule the file for deletion on next reboot
   162→    // This is necessary because the executable is currently running
   163→    unsafe {
   164→        let result = MoveFileExW(
   165→            PCWSTR(exe_wide.as_ptr()),
   166→            PCWSTR::null(),
   167→            MOVEFILE_DELAY_UNTIL_REBOOT,
   168→        );
   169→
   170→        if result.is_err() {
   171→            warn!("Failed to schedule file deletion: {:?}. File may need manual removal.", result);
   172→        } else {
   173→            info!("Executable scheduled for deletion on reboot");
   174→        }
   175→    }
   176→
   177→    Ok(())
   178→}
   179→
   180→#[cfg(not(windows))]
   181→pub fn add_to_startup() -> Result<()> {
   182→    warn!("Startup persistence not implemented for this platform");
   183→    Ok(())
   184→}
   185→
   186→#[cfg(not(windows))]
   187→pub fn remove_from_startup() -> Result<()> {
   188→    Ok(())
   189→}
   190→
   191→#[cfg(not(windows))]
   192→pub fn uninstall() -> Result<()> {
   193→    warn!("Uninstall not implemented for this platform");
   194→    Ok(())
   195→}
   196→

/// Full uninstall: remove from startup and delete the executable
#[cfg(windows)]
pub fn uninstall() -> Result<()> {
    use std::ffi::OsStr;
    use std::os::windows::ffi::OsStrExt;
    use windows::Win32::Storage::FileSystem::{MoveFileExW, MOVEFILE_DELAY_UNTIL_REBOOT};

    info!("Uninstalling agent");

    // First remove from startup
    let _ = remove_from_startup();

    // Get the path to the current executable
    let exe_path = std::env::current_exe()?;
    let exe_path_str = exe_path.to_string_lossy();

    info!("Scheduling deletion of: {}", exe_path_str);

    // Convert path to wide string
    let exe_wide: Vec<u16> = OsStr::new(&*exe_path_str)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    // Schedule the file for deletion on next reboot
    // This is necessary because the executable is currently running
    unsafe {
        let result = MoveFileExW(
            PCWSTR(exe_wide.as_ptr()),
            PCWSTR::null(),
            MOVEFILE_DELAY_UNTIL_REBOOT,
        );

        if result.is_err() {
            warn!("Failed to schedule file deletion: {:?}. File may need manual removal.", result);
        } else {
            info!("Executable scheduled for deletion on reboot");
        }
    }

    Ok(())
}

#[cfg(not(windows))]
pub fn add_to_startup() -> Result<()> {
    warn!("Startup persistence not implemented for this platform");
    Ok(())
}

#[cfg(not(windows))]
pub fn remove_from_startup() -> Result<()> {
    Ok(())
}

#[cfg(not(windows))]
pub fn uninstall() -> Result<()> {
    warn!("Uninstall not implemented for this platform");
    Ok(())
}

//! Startup persistence for the agent
//!
//! Handles adding/removing the agent from Windows startup.

use anyhow::Result;
use tracing::{info, warn, error};

#[cfg(windows)]
use windows::Win32::System::Registry::{
    RegOpenKeyExW, RegSetValueExW, RegDeleteValueW, RegCloseKey,
    HKEY_CURRENT_USER, KEY_WRITE, REG_SZ,
};
#[cfg(windows)]
use windows::core::PCWSTR;

const STARTUP_KEY: &str = r"Software\Microsoft\Windows\CurrentVersion\Run";
const STARTUP_VALUE_NAME: &str = "GuruConnect";

/// Add the current executable to Windows startup
#[cfg(windows)]
pub fn add_to_startup() -> Result<()> {
    use std::ffi::OsStr;
    use std::os::windows::ffi::OsStrExt;

    // Get the path to the current executable
    let exe_path = std::env::current_exe()?;
    let exe_path_str = exe_path.to_string_lossy();

    info!("Adding to startup: {}", exe_path_str);

    // Convert strings to wide strings
    let key_path: Vec<u16> = OsStr::new(STARTUP_KEY)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();
    let value_name: Vec<u16> = OsStr::new(STARTUP_VALUE_NAME)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();
    let value_data: Vec<u16> = OsStr::new(&*exe_path_str)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    unsafe {
        let mut hkey = windows::Win32::Foundation::HANDLE::default();

        // Open the Run key
        let result = RegOpenKeyExW(
            HKEY_CURRENT_USER,
            PCWSTR(key_path.as_ptr()),
            0,
            KEY_WRITE,
            &mut hkey as *mut _ as *mut _,
        );

        if result.is_err() {
            anyhow::bail!("Failed to open registry key: {:?}", result);
        }

        let hkey_raw = std::mem::transmute::<_, windows::Win32::System::Registry::HKEY>(hkey);

        // Set the value
        let data_bytes = std::slice::from_raw_parts(
            value_data.as_ptr() as *const u8,
            value_data.len() * 2,
        );

        let set_result = RegSetValueExW(
            hkey_raw,
            PCWSTR(value_name.as_ptr()),
            0,
            REG_SZ,
            Some(data_bytes),
        );

        let _ = RegCloseKey(hkey_raw);

        if set_result.is_err() {
            anyhow::bail!("Failed to set registry value: {:?}", set_result);
        }
    }

    info!("Successfully added to startup");
    Ok(())
}

/// Remove the agent from Windows startup
#[cfg(windows)]
pub fn remove_from_startup() -> Result<()> {
    use std::ffi::OsStr;
    use std::os::windows::ffi::OsStrExt;

    info!("Removing from startup");

    let key_path: Vec<u16> = OsStr::new(STARTUP_KEY)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();
    let value_name: Vec<u16> = OsStr::new(STARTUP_VALUE_NAME)
        .encode_wide()
        .chain(std::iter::once(0))
        .collect();

    unsafe {
        let mut hkey = windows::Win32::Foundation::HANDLE::default();

        let result = RegOpenKeyExW(
            HKEY_CURRENT_USER,
            PCWSTR(key_path.as_ptr()),
            0,
            KEY_WRITE,
            &mut hkey as *mut _ as *mut _,
        );

        if result.is_err() {
            warn!("Failed to open registry key for removal: {:?}", result);
            return Ok(()); // Not an error if key doesn't exist
        }

        let hkey_raw = std::mem::transmute::<_, windows::Win32::System::Registry::HKEY>(hkey);

        let delete_result = RegDeleteValueW(hkey_raw, PCWSTR(value_name.as_ptr()));

        let _ = RegCloseKey(hkey_raw);

        if delete_result.is_err() {
            warn!("Registry value may not exist: {:?}", delete_result);
        } else {
            info!("Successfully removed from startup");
        }
    }

    Ok(())
}

#[cfg(not(windows))]
pub fn add_to_startup() -> Result<()> {
    warn!("Startup persistence not implemented for this platform");
    Ok(())
}

#[cfg(not(windows))]
pub fn remove_from_startup() -> Result<()> {
    Ok(())
}

+/// Full uninstall: remove from startup and delete the executable

+    use std::os::windows::ffi::OsStrExt;

+    use windows::Win32::Storage::FileSystem::{MoveFileExW, MOVEFILE_DELAY_UNTIL_REBOOT};

+    // Get the path to the current executable

+    let exe_path = std::env::current_exe()?;

+    let exe_path_str = exe_path.to_string_lossy();

+    info!("Scheduling deletion of: {}", exe_path_str);

+    let exe_wide: Vec<u16> = OsStr::new(&*exe_path_str)

+    // Schedule the file for deletion on next reboot

+    // This is necessary because the executable is currently running

+            MOVEFILE_DELAY_UNTIL_REBOOT,

+            warn!("Failed to schedule file deletion: {:?}. File may need manual removal.", result);

+            info!("Executable scheduled for deletion on reboot");

+    warn!("Uninstall not implemented for this platform");


Now I need to update the agent's session handling to call uninstall when it receives a disconnect message from an administrator. Let me check the session/mod.rs to see how disconnects are handled.


Now let me update the session module to handle admin disconnects with uninstall:


            Some(message::Payload::Disconnect(disc)) => {
                tracing::info!("Disconnect requested: {}", disc.reason);
                // Check if this is a cancellation
                if disc.reason.contains("cancelled") {
                    return Err(anyhow::anyhow!("SESSION_CANCELLED: {}", disc.reason));
                }
                return Err(anyhow::anyhow!("Disconnect: {}", disc.reason));
            }

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