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:
    58→
    59→    ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, db, agent_id, agent_name, support_code))
    60→}
    61→
    62→/// WebSocket handler for viewer connections
    63→pub async fn viewer_ws_handler(
    64→    ws: WebSocketUpgrade,
    65→    State(state): State<AppState>,
    66→    Query(params): Query<ViewerParams>,
    67→) -> impl IntoResponse {
    68→    let session_id = params.session_id;
    69→    let viewer_name = params.viewer_name;
    70→    let sessions = state.sessions.clone();
    71→    let db = state.db.clone();
    72→
    73→    ws.on_upgrade(move |socket| handle_viewer_connection(socket, sessions, db, session_id, viewer_name))
    74→}
    75→
    76→/// Handle an agent WebSocket connection
    77→async fn handle_agent_connection(
    78→    socket: WebSocket,

/// 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 viewer_name = params.viewer_name;
    let sessions = state.sessions.clone();
    let db = state.db.clone();

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

//! 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::db::{self, Database};
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,
    #[serde(default = "default_viewer_name")]
    viewer_name: String,
}

fn default_viewer_name() -> String {
    "Technician".to_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();
    let db = state.db.clone();

    ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, db, 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 viewer_name = params.viewer_name;
    let sessions = state.sessions.clone();

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

/// 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
    // Persistent agents (no support code) keep their session when disconnected
    let is_persistent = support_code.is_none();
    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone(), is_persistent).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();
    // Mark agent as disconnected (persistent agents stay in list as offline)
    sessions_cleanup.mark_agent_disconnected(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,
    viewer_name: 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(), viewer_name.clone()).await {
        Some(channels) => channels,
        None => {
            warn!("Session not found: {}", session_id);
            return;
        }
    };

    info!("Viewer {} ({}) joined session: {}", viewer_name, 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);
}

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

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


Now I need to update the handle_agent_connection function signature and add database logging. Let me find and update that function.


Edit

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

/// Handle an agent WebSocket connection
async fn handle_agent_connection(
    socket: WebSocket,
    sessions: SessionManager,
    support_codes: crate::support_codes::SupportCodeManager,
    db: Option<Database>,
    agent_id: String,
    agent_name: String,
    support_code: Option<String>,
) {


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:
    72→
    73→    ws.on_upgrade(move |socket| handle_viewer_connection(socket, sessions, db, session_id, viewer_name))
    74→}
    75→
    76→/// Handle an agent WebSocket connection
    77→async fn handle_agent_connection(
    78→    socket: WebSocket,
    79→    sessions: SessionManager,
    80→    support_codes: crate::support_codes::SupportCodeManager,
    81→    db: Option<Database>,
    82→    agent_id: String,
    83→    agent_name: String,
    84→    support_code: Option<String>,
    85→) {
    86→    info!("Agent connected: {} ({})", agent_name, agent_id);
    87→
    88→    let (mut ws_sender, mut ws_receiver) = socket.split();
    89→

/// Handle an agent WebSocket connection
async fn handle_agent_connection(
    socket: WebSocket,
    sessions: SessionManager,
    support_codes: crate::support_codes::SupportCodeManager,
    db: Option<Database>,
    agent_id: String,
    agent_name: String,
    support_code: Option<String>,
) {

//! 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::db::{self, Database};
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,
    #[serde(default = "default_viewer_name")]
    viewer_name: String,
}

fn default_viewer_name() -> String {
    "Technician".to_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();
    let db = state.db.clone();

    ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, db, 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 viewer_name = params.viewer_name;
    let sessions = state.sessions.clone();
    let db = state.db.clone();

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

/// 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
    // Persistent agents (no support code) keep their session when disconnected
    let is_persistent = support_code.is_none();
    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone(), is_persistent).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();
    // Mark agent as disconnected (persistent agents stay in list as offline)
    sessions_cleanup.mark_agent_disconnected(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,
    viewer_name: 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(), viewer_name.clone()).await {
        Some(channels) => channels,
        None => {
            warn!("Session not found: {}", session_id);
            return;
        }
    };

    info!("Viewer {} ({}) joined session: {}", viewer_name, 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);
}


Good. Now I need to add the database logging after the agent registers. Let me find the right place in the file.


   100→            };
   101→            let mut buf = Vec::new();
   102→            if prost::Message::encode(&disconnect_msg, &mut buf).is_ok() {
   103→                let _ = ws_sender.send(Message::Binary(buf.into())).await;
   104→            }
   105→            let _ = ws_sender.close().await;
   106→            return;
   107→        }
   108→    }
   109→
   110→    // Register the agent and get channels
   111→    // Persistent agents (no support code) keep their session when disconnected
   112→    let is_persistent = support_code.is_none();
   113→    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone(), is_persistent).await;
   114→
   115→    info!("Session created: {} (agent in idle mode)", session_id);
   116→
   117→    // If a support code was provided, mark it as connected
   118→    if let Some(ref code) = support_code {
   119→        info!("Linking support code {} to session {}", code, session_id);
   120→        support_codes.mark_connected(code, Some(agent_name.clone()), Some(agent_id.clone())).await;
   121→        support_codes.link_session(code, session_id).await;
   122→    }
   123→
   124→    // Use Arc<Mutex> for sender so we can use it from multiple places
   125→    let ws_sender = std::sync::Arc::new(tokio::sync::Mutex::new(ws_sender));
   126→    let ws_sender_input = ws_sender.clone();
   127→    let ws_sender_cancel = ws_sender.clone();
   128→
   129→    // Task to forward input events from viewers to agent
   130→    let input_forward = tokio::spawn(async move {
   131→        while let Some(input_data) = input_rx.recv().await {
   132→            let mut sender = ws_sender_input.lock().await;
   133→            if sender.send(Message::Binary(input_data.into())).await.is_err() {
   134→                break;
   135→            }
   136→        }
   137→    });
   138→
   139→    let sessions_cleanup = sessions.clone();
   140→    let sessions_status = sessions.clone();
   141→    let support_codes_cleanup = support_codes.clone();
   142→    let support_code_cleanup = support_code.clone();
   143→    let support_code_check = support_code.clone();
   144→    let support_codes_check = support_codes.clone();
   145→
   146→    // Task to check for cancellation every 2 seconds
   147→    let cancel_check = tokio::spawn(async move {
   148→        let mut interval = tokio::time::interval(std::time::Duration::from_secs(2));
   149→        loop {
   150→            interval.tick().await;
   151→            if let Some(ref code) = support_code_check {
   152→                if support_codes_check.is_cancelled(code).await {
   153→                    info!("Support code {} was cancelled, disconnecting agent", code);
   154→                    // Send disconnect message
   155→                    let disconnect_msg = proto::Message {
   156→                        payload: Some(proto::message::Payload::Disconnect(proto::Disconnect {
   157→                            reason: "Support session was cancelled by technician".to_string(),
   158→                        })),
   159→                    };

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

            };
            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
    // Persistent agents (no support code) keep their session when disconnected
    let is_persistent = support_code.is_none();
    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone(), is_persistent).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(),
                        })),
                    };


I need to add database operations after session registration. Let me add the database insert after line 115.


    // Register the agent and get channels
    // Persistent agents (no support code) keep their session when disconnected
    let is_persistent = support_code.is_none();
    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone(), is_persistent).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;
    }

    // Register the agent and get channels
    // Persistent agents (no support code) keep their session when disconnected
    let is_persistent = support_code.is_none();
    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone(), is_persistent).await;

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

    // Database: upsert machine and create session record
    let machine_id = if let Some(ref db) = db {
        match db::machines::upsert_machine(db.pool(), &agent_id, &agent_name, is_persistent).await {
            Ok(machine) => {
                // Create session record
                let _ = db::sessions::create_session(
                    db.pool(),
                    session_id,
                    machine.id,
                    support_code.is_some(),
                    support_code.as_deref(),
                ).await;

                // Log session started event
                let _ = db::events::log_event(
                    db.pool(),
                    session_id,
                    db::events::EventTypes::SESSION_STARTED,
                    None, None, None, None,
                ).await;

                Some(machine.id)
            }
            Err(e) => {
                warn!("Failed to upsert machine in database: {}", e);
                None
            }
        }
    } else {
        None
    };

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

        // Database: update support code
        if let Some(ref db) = db {
            let _ = db::support_codes::mark_code_connected(
                db.pool(),
                code,
                Some(session_id),
                Some(&agent_name),
                Some(&agent_id),
            ).await;
        }
    }


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:
   106→            return;
   107→        }
   108→    }
   109→
   110→    // Register the agent and get channels
   111→    // Persistent agents (no support code) keep their session when disconnected
   112→    let is_persistent = support_code.is_none();
   113→    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone(), is_persistent).await;
   114→
   115→    info!("Session created: {} (agent in idle mode)", session_id);
   116→
   117→    // Database: upsert machine and create session record
   118→    let machine_id = if let Some(ref db) = db {
   119→        match db::machines::upsert_machine(db.pool(), &agent_id, &agent_name, is_persistent).await {
   120→            Ok(machine) => {
   121→                // Create session record
   122→                let _ = db::sessions::create_session(
   123→                    db.pool(),
   124→                    session_id,
   125→                    machine.id,
   126→                    support_code.is_some(),
   127→                    support_code.as_deref(),
   128→                ).await;
   129→
   130→                // Log session started event
   131→                let _ = db::events::log_event(
   132→                    db.pool(),
   133→                    session_id,
   134→                    db::events::EventTypes::SESSION_STARTED,
   135→                    None, None, None, None,
   136→                ).await;
   137→
   138→                Some(machine.id)
   139→            }
   140→            Err(e) => {
   141→                warn!("Failed to upsert machine in database: {}", e);
   142→                None
   143→            }
   144→        }
   145→    } else {
   146→        None
   147→    };
   148→
   149→    // If a support code was provided, mark it as connected
   150→    if let Some(ref code) = support_code {
   151→        info!("Linking support code {} to session {}", code, session_id);
   152→        support_codes.mark_connected(code, Some(agent_name.clone()), Some(agent_id.clone())).await;
   153→        support_codes.link_session(code, session_id).await;
   154→
   155→        // Database: update support code
   156→        if let Some(ref db) = db {
   157→            let _ = db::support_codes::mark_code_connected(
   158→                db.pool(),
   159→                code,
   160→                Some(session_id),
   161→                Some(&agent_name),
   162→                Some(&agent_id),
   163→            ).await;
   164→        }
   165→    }
   166→
   167→    // Use Arc<Mutex> for sender so we can use it from multiple places
   168→    let ws_sender = std::sync::Arc::new(tokio::sync::Mutex::new(ws_sender));
   169→    let ws_sender_input = ws_sender.clone();

    // Register the agent and get channels
    // Persistent agents (no support code) keep their session when disconnected
    let is_persistent = support_code.is_none();
    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone(), is_persistent).await;

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

    // Database: upsert machine and create session record
    let machine_id = if let Some(ref db) = db {
        match db::machines::upsert_machine(db.pool(), &agent_id, &agent_name, is_persistent).await {
            Ok(machine) => {
                // Create session record
                let _ = db::sessions::create_session(
                    db.pool(),
                    session_id,
                    machine.id,
                    support_code.is_some(),
                    support_code.as_deref(),
                ).await;

                // Log session started event
                let _ = db::events::log_event(
                    db.pool(),
                    session_id,
                    db::events::EventTypes::SESSION_STARTED,
                    None, None, None, None,
                ).await;

                Some(machine.id)
            }
            Err(e) => {
                warn!("Failed to upsert machine in database: {}", e);
                None
            }
        }
    } else {
        None
    };

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

        // Database: update support code
        if let Some(ref db) = db {
            let _ = db::support_codes::mark_code_connected(
                db.pool(),
                code,
                Some(session_id),
                Some(&agent_name),
                Some(&agent_id),
            ).await;
        }
    }

//! 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::db::{self, Database};
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,
    #[serde(default = "default_viewer_name")]
    viewer_name: String,
}

fn default_viewer_name() -> String {
    "Technician".to_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();
    let db = state.db.clone();

    ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, db, 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 viewer_name = params.viewer_name;
    let sessions = state.sessions.clone();
    let db = state.db.clone();

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

/// Handle an agent WebSocket connection
async fn handle_agent_connection(
    socket: WebSocket,
    sessions: SessionManager,
    support_codes: crate::support_codes::SupportCodeManager,
    db: Option<Database>,
    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
    // Persistent agents (no support code) keep their session when disconnected
    let is_persistent = support_code.is_none();
    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone(), is_persistent).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();
    // Mark agent as disconnected (persistent agents stay in list as offline)
    sessions_cleanup.mark_agent_disconnected(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,
    viewer_name: 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(), viewer_name.clone()).await {
        Some(channels) => channels,
        None => {
            warn!("Session not found: {}", session_id);
            return;
        }
    };

    info!("Viewer {} ({}) joined session: {}", viewer_name, 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);
}

+    // Database: upsert machine and create session record

+    let machine_id = if let Some(ref db) = db {

+        match db::machines::upsert_machine(db.pool(), &agent_id, &agent_name, is_persistent).await {

+                // Create session record

+                let _ = db::sessions::create_session(

+                    support_code.is_some(),

+                    support_code.as_deref(),

+                // Log session started event

+                let _ = db::events::log_event(

+                    db::events::EventTypes::SESSION_STARTED,

+                    None, None, None, None,

+                warn!("Failed to upsert machine in database: {}", e);

+        // Database: update support code

+            let _ = db::support_codes::mark_code_connected(


Conversation compacted


This session is being continued from a previous conversation that ran out of context. The conversation is summarized below:
Analysis:
Let me chronologically analyze the conversation:

1. **Initial Context Recovery**: The conversation started with a summary from a previous session about GuruConnect - implementing connected technician tracking (showing "Mike Connected (3 min)" on machine bars).

2. **Completing Technician Tracking Feature**:
   - Updated dashboard.html with formatViewers() and formatDuration() functions
   - Updated viewer.html to pass viewer_name parameter in WebSocket URL
   - Updated dashboard chat connection to also pass viewer_name
   - Deployed static files and server binary to build server

3. **Deployment Challenges**:
   - Had to find correct paths on build server (/home/guru/guru-connect/server/static/)
   - Built server with Linux target: `cargo build --release --target x86_64-unknown-linux-gnu -p guruconnect-server`
   - Used pkill workaround to restart service since we don't have passwordless sudo
   - Successfully deployed by: rename old binary → copy new → pkill old process → systemd auto-restarts

4. **Updated credentials.md** with proper deployment method documentation

5. **User's New Request - Database Persistence**:
   - User asked: "Machine persistence needs to survive reboots too. Also we need to track all interactions with machines, do we need to setup a database for that?"
   - User chose PostgreSQL over SQLite

6. **Plan Mode - Database Integration Design**:
   - Explored GuruRMM's PostgreSQL patterns (sqlx 0.8, connection pooling, migrations)
   - Explored GuruConnect's current in-memory architecture
   - Designed schema with 4 tables: connect_machines, connect_sessions, connect_session_events, connect_support_codes
   - Created comprehensive plan with 8 implementation steps

7. **Implementation Started**:
   - Created database via RMM agent commands
   - Created migration file (001_initial_schema.sql)
   - Created db modules: mod.rs, machines.rs, sessions.rs, events.rs, support_codes.rs
   - Updated config.rs with database_max_connections
   - Updated main.rs with database initialization and machine restoration
   - Added restore_offline_machine() to SessionManager
   - Started updating relay/mod.rs with database logging

8. **Current State**:
   - Partially through updating relay/mod.rs
   - Added db parameter to handler functions
   - Added database logging after agent registration (upsert_machine, create_session, log_event)
   - Still need to add cleanup logging on disconnect and viewer join/leave logging

Summary:
1. Primary Request and Intent:
   - **Machine Persistence**: User wants agents/machines to survive server reboots - show as "offline" until they reconnect
   - **Audit Logging**: Track all technician interactions - who connected, when, session durations
   - **Database Choice**: User explicitly chose PostgreSQL (use existing instance on build server 172.16.3.30)
   - Earlier in session: Completing technician tracking feature showing "Mike Connected (3 min)" on dashboard

2. Key Technical Concepts:
   - PostgreSQL with sqlx 0.8 for async database operations
   - Connection pooling with PgPoolOptions
   - sqlx::migrate!() macro for automatic migrations
   - Hybrid storage: in-memory for real-time performance, database for persistence
   - Tables prefixed with `connect_*` to coexist with GuruRMM
   - Fire-and-forget logging (DB writes don't block WebSocket handling)
   - Optional database - server works without DATABASE_URL set
   - Service restart via pkill (services run as guru user, no sudo needed)

3. Files and Code Sections:

   - **server/migrations/001_initial_schema.sql** (CREATED)
     - Initial database schema with 4 tables
     ```sql
     CREATE TABLE connect_machines (
         id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
         agent_id VARCHAR(255) UNIQUE NOT NULL,
         hostname VARCHAR(255) NOT NULL,
         os_version VARCHAR(255),
         is_elevated BOOLEAN DEFAULT FALSE,
         is_persistent BOOLEAN DEFAULT TRUE,
         first_seen TIMESTAMPTZ DEFAULT NOW(),
         last_seen TIMESTAMPTZ DEFAULT NOW(),
         last_session_id UUID,
         status VARCHAR(20) DEFAULT 'offline',
         created_at TIMESTAMPTZ DEFAULT NOW(),
         updated_at TIMESTAMPTZ DEFAULT NOW()
     );
     -- Plus connect_sessions, connect_session_events, connect_support_codes tables
     ```

   - **server/src/db/mod.rs** (REWRITTEN)
     - Database connection pool wrapper with connect/migrate methods
     ```rust
     #[derive(Clone)]
     pub struct Database {
         pool: PgPool,
     }

     impl Database {
         pub async fn connect(database_url: &str, max_connections: u32) -> Result<Self>
         pub async fn migrate(&self) -> Result<()>
         pub fn pool(&self) -> &PgPool
     }
     ```

   - **server/src/db/machines.rs** (CREATED)
     - Machine CRUD operations: upsert_machine, update_machine_status, get_all_machines, mark_machine_offline

   - **server/src/db/sessions.rs** (CREATED)
     - Session logging: create_session, end_session, get_session

   - **server/src/db/events.rs** (CREATED)
     - Audit event logging with EventTypes constants (SESSION_STARTED, VIEWER_JOINED, etc.)
     ```rust
     pub async fn log_event(
         pool: &PgPool,
         session_id: Uuid,
         event_type: &str,
         viewer_id: Option<&str>,
         viewer_name: Option<&str>,
         details: Option<JsonValue>,
         ip_address: Option<IpAddr>,
     ) -> Result<i64, sqlx::Error>
     ```

   - **server/src/db/support_codes.rs** (CREATED)
     - Support code persistence: create_support_code, mark_code_connected, mark_code_completed

   - **server/src/config.rs** (UPDATED)
     - Added database_max_connections field (default: 5)

   - **server/src/main.rs** (UPDATED)
     - Added `db: Option<db::Database>` to AppState
     - Added database initialization with migration running
     - Added machine restoration from database on startup
     ```rust
     // Restore persistent machines from database
     if let Some(ref db) = database {
         match db::machines::get_all_machines(db.pool()).await {
             Ok(machines) => {
                 for machine in machines {
                     sessions.restore_offline_machine(&machine.agent_id, &machine.hostname).await;
                 }
             }
             // ...
         }
     }
     ```

   - **server/src/session/mod.rs** (UPDATED)
     - Added restore_offline_machine() method for startup recovery
     ```rust
     pub async fn restore_offline_machine(&self, agent_id: &str, hostname: &str) -> SessionId
     ```

   - **server/src/relay/mod.rs** (PARTIALLY UPDATED)
     - Added `db: Option<Database>` parameter to handlers
     - Added database logging after agent registration
     ```rust
     // Database: upsert machine and create session record
     let machine_id = if let Some(ref db) = db {
         match db::machines::upsert_machine(db.pool(), &agent_id, &agent_name, is_persistent).await {
             Ok(machine) => {
                 let _ = db::sessions::create_session(...).await;
                 let _ = db::events::log_event(..., db::events::EventTypes::SESSION_STARTED, ...).await;
                 Some(machine.id)
             }
             // ...
         }
     };
     ```

4. Errors and fixes:
   - **Service restart without sudo**: Used pkill method since services run as guru user
   - **Binary busy during copy**: Renamed old binary first, then copied new, then killed old process
   - **Build target mismatch**: Had to specify `--target x86_64-unknown-linux-gnu` for Linux build
   - **Database creation via SSH**: Used RMM agent to run sudo commands since SSH sudo doesn't work with special char password

5. Problem Solving:
   - Solved deployment by documenting pkill method in credentials.md
   - Created database via RMM agent (runs as root) since SSH sudo has password issues
   - Designed hybrid in-memory + database architecture for performance

6. All user messages:
   - "Machine persistence needs to survive reboots too. Also we need to track all interactions with machines, do we need to setup a database for that?"
   - User selected "PostgreSQL (Recommended)" when asked about database choice
   - User approved the database integration plan

7. Pending Tasks:
   - Complete relay/mod.rs updates (add disconnect/cleanup logging, viewer join/leave events)
   - Update handle_viewer_connection with database parameter and logging
   - Add cleanup logging on agent disconnect (end_session, mark_machine_offline, log SESSION_ENDED)
   - Deploy and test the complete implementation
   - Verify database was created successfully via RMM

8. Current Work:
   Implementing PostgreSQL database integration for GuruConnect. Was in the middle of updating **server/src/relay/mod.rs** to add database logging. Completed:
   - Added db parameter to handler functions
   - Added database logging after agent registration (upsert_machine, create_session, log_event for SESSION_STARTED)
   - Added support code database update when code is used

   Still need to add in relay/mod.rs:
   - Cleanup logging on agent disconnect
   - handle_viewer_connection updates (db parameter, viewer_joined/viewer_left events)

9. Optional Next Step:
   Continue updating relay/mod.rs to add:
   1. Database cleanup logging on agent disconnect (end_session, mark_machine_offline, log SESSION_ENDED event)
   2. Update handle_viewer_connection signature with db parameter
   3. Add viewer_joined and viewer_left event logging

   From the current work: Was editing relay/mod.rs after successfully adding the agent connection database logging. The agent_ws_handler and handle_agent_connection now have db parameter and initial logging. Need to add the cleanup logging and viewer handler updates..
Please continue the conversation from where we left it off without asking the user any further questions. Continue with the last task that you were asked to work on.


The user wants me to continue where we left off. Based on the summary, I was in the middle of updating `relay/mod.rs` to add database logging. I've already:
1. Added db parameter to handler functions
2. Added database logging after agent registration (upsert_machine, create_session, log_event for SESSION_STARTED)
3. Added support code database update when code is used

I still need to:
1. Add cleanup logging on agent disconnect (end_session, mark_machine_offline, log SESSION_ENDED event)
2. Update handle_viewer_connection signature with db parameter
3. Add viewer_joined and viewer_left event logging

Let me look at the current state of relay/mod.rs to see exactly what's there and what needs to be added.


Continuing with the database integration. I need to complete the relay/mod.rs updates by adding:
1. Cleanup logging on agent disconnect
2. Viewer connection database logging

Let me check the current state of relay/mod.rs to see what's already done and what remains:


   260→            }
   261→            Ok(Message::Close(_)) => {
   262→                info!("Agent disconnected: {}", agent_id);
   263→                break;
   264→            }
   265→            Ok(Message::Ping(data)) => {
   266→                // Pong is handled automatically by axum
   267→                let _ = data;
   268→            }
   269→            Ok(_) => {}
   270→            Err(e) => {
   271→                error!("WebSocket error from agent {}: {}", agent_id, e);
   272→                break;
   273→            }
   274→        }
   275→    }
   276→
   277→    // Cleanup
   278→    input_forward.abort();
   279→    cancel_check.abort();
   280→    // Mark agent as disconnected (persistent agents stay in list as offline)
   281→    sessions_cleanup.mark_agent_disconnected(session_id).await;
   282→
   283→    // Mark support code as completed if one was used (unless cancelled)
   284→    if let Some(ref code) = support_code_cleanup {
   285→        if !support_codes_cleanup.is_cancelled(code).await {
   286→            support_codes_cleanup.mark_completed(code).await;
   287→            info!("Support code {} marked as completed", code);
   288→        }
   289→    }
   290→
   291→    info!("Session {} ended", session_id);
   292→}
   293→
   294→/// Handle a viewer WebSocket connection
   295→async fn handle_viewer_connection(
   296→    socket: WebSocket,
   297→    sessions: SessionManager,
   298→    session_id_str: String,
   299→    viewer_name: String,
   300→) {
   301→    // Parse session ID
   302→    let session_id = match uuid::Uuid::parse_str(&session_id_str) {
   303→        Ok(id) => id,
   304→        Err(_) => {
   305→            warn!("Invalid session ID: {}", session_id_str);
   306→            return;
   307→        }
   308→    };
   309→
   310→    // Generate unique viewer ID
   311→    let viewer_id = Uuid::new_v4().to_string();
   312→
   313→    // Join the session (this sends StartStream to agent if first viewer)
   314→    let (mut frame_rx, input_tx) = match sessions.join_session(session_id, viewer_id.clone(), viewer_name.clone()).await {
   315→        Some(channels) => channels,
   316→        None => {
   317→            warn!("Session not found: {}", session_id);
   318→            return;
   319→        }
   320→    };
   321→
   322→    info!("Viewer {} ({}) joined session: {}", viewer_name, viewer_id, session_id);
   323→
   324→    let (mut ws_sender, mut ws_receiver) = socket.split();
   325→
   326→    // Task to forward frames from agent to this viewer
   327→    let frame_forward = tokio::spawn(async move {
   328→        while let Ok(frame_data) = frame_rx.recv().await {
   329→            if ws_sender.send(Message::Binary(frame_data.into())).await.is_err() {
   330→                break;
   331→            }
   332→        }
   333→    });
   334→
   335→    let sessions_cleanup = sessions.clone();
   336→    let viewer_id_cleanup = viewer_id.clone();
   337→
   338→    // Main loop: receive input from viewer and forward to agent
   339→    while let Some(msg) = ws_receiver.next().await {
   340→        match msg {
   341→            Ok(Message::Binary(data)) => {
   342→                // Try to decode as protobuf message
   343→                match proto::Message::decode(data.as_ref()) {
   344→                    Ok(proto_msg) => {
   345→                        match &proto_msg.payload {
   346→                            Some(proto::message::Payload::MouseEvent(_)) |
   347→                            Some(proto::message::Payload::KeyEvent(_)) |
   348→                            Some(proto::message::Payload::SpecialKey(_)) => {
   349→                                // Forward input to agent
   350→                                let _ = input_tx.send(data.to_vec()).await;
   351→                            }
   352→                            Some(proto::message::Payload::ChatMessage(chat)) => {
   353→                                // Forward chat message to agent
   354→                                info!("Chat from technician: {}", chat.content);
   355→                                let _ = input_tx.send(data.to_vec()).await;
   356→                            }
   357→                            _ => {}
   358→                        }
   359→                    }
   360→                    Err(e) => {
   361→                        warn!("Failed to decode viewer message: {}", e);
   362→                    }
   363→                }
   364→            }
   365→            Ok(Message::Close(_)) => {
   366→                info!("Viewer {} disconnected from session: {}", viewer_id, session_id);
   367→                break;
   368→            }
   369→            Ok(_) => {}
   370→            Err(e) => {
   371→                error!("WebSocket error from viewer {}: {}", viewer_id, e);
   372→                break;
   373→            }
   374→        }
   375→    }
   376→
   377→    // Cleanup (this sends StopStream to agent if last viewer)
   378→    frame_forward.abort();
   379→    sessions_cleanup.leave_session(session_id, &viewer_id_cleanup).await;
   380→    info!("Viewer {} left session: {}", viewer_id_cleanup, session_id);
   381→}
   382→

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

            }
            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();
    // Mark agent as disconnected (persistent agents stay in list as offline)
    sessions_cleanup.mark_agent_disconnected(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,
    viewer_name: 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(), viewer_name.clone()).await {
        Some(channels) => channels,
        None => {
            warn!("Session not found: {}", session_id);
            return;
        }
    };

    info!("Viewer {} ({}) joined session: {}", viewer_name, 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);
}


I can see the current state. The handle_viewer_connection function doesn't have the `db` parameter yet, and I need to:

1. Add agent disconnect database logging in handle_agent_connection (after line 291)
2. Update handle_viewer_connection to accept db parameter and add viewer_joined/viewer_left event logging

Let me update the relay/mod.rs file. First I'll add the cleanup logging for agent disconnect, then update the viewer handler.


I can see the current state. I need to:
1. Add database cleanup logging on agent disconnect (after line 281)
2. Update handle_viewer_connection to accept the db parameter and add event logging

Let me update the cleanup section first:


    // Cleanup
    input_forward.abort();
    cancel_check.abort();
    // Mark agent as disconnected (persistent agents stay in list as offline)
    sessions_cleanup.mark_agent_disconnected(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);

    // Cleanup
    input_forward.abort();
    cancel_check.abort();
    // Mark agent as disconnected (persistent agents stay in list as offline)
    sessions_cleanup.mark_agent_disconnected(session_id).await;

    // Database: end session and mark machine offline
    if let Some(ref db) = db {
        // End the session record
        let _ = db::sessions::end_session(db.pool(), session_id).await;

        // Mark machine as offline
        if let Some(mid) = machine_id {
            let _ = db::machines::mark_machine_offline(db.pool(), mid).await;
        }

        // Log session ended event
        let _ = db::events::log_event(
            db.pool(),
            session_id,
            db::events::EventTypes::SESSION_ENDED,
            None, None, None, None,
        ).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;

            // Database: mark code as completed
            if let Some(ref db) = db {
                let _ = db::support_codes::mark_code_completed(db.pool(), code).await;
            }

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

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


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:
   273→            }
   274→        }
   275→    }
   276→
   277→    // Cleanup
   278→    input_forward.abort();
   279→    cancel_check.abort();
   280→    // Mark agent as disconnected (persistent agents stay in list as offline)
   281→    sessions_cleanup.mark_agent_disconnected(session_id).await;
   282→
   283→    // Database: end session and mark machine offline
   284→    if let Some(ref db) = db {
   285→        // End the session record
   286→        let _ = db::sessions::end_session(db.pool(), session_id).await;
   287→
   288→        // Mark machine as offline
   289→        if let Some(mid) = machine_id {
   290→            let _ = db::machines::mark_machine_offline(db.pool(), mid).await;
   291→        }
   292→
   293→        // Log session ended event
   294→        let _ = db::events::log_event(
   295→            db.pool(),
   296→            session_id,
   297→            db::events::EventTypes::SESSION_ENDED,
   298→            None, None, None, None,
   299→        ).await;
   300→    }
   301→
   302→    // Mark support code as completed if one was used (unless cancelled)
   303→    if let Some(ref code) = support_code_cleanup {
   304→        if !support_codes_cleanup.is_cancelled(code).await {
   305→            support_codes_cleanup.mark_completed(code).await;
   306→
   307→            // Database: mark code as completed
   308→            if let Some(ref db) = db {
   309→                let _ = db::support_codes::mark_code_completed(db.pool(), code).await;
   310→            }
   311→
   312→            info!("Support code {} marked as completed", code);
   313→        }
   314→    }
   315→
   316→    info!("Session {} ended", session_id);
   317→}
   318→
   319→/// Handle a viewer WebSocket connection
   320→async fn handle_viewer_connection(

    // Cleanup
    input_forward.abort();
    cancel_check.abort();
    // Mark agent as disconnected (persistent agents stay in list as offline)
    sessions_cleanup.mark_agent_disconnected(session_id).await;

    // Database: end session and mark machine offline
    if let Some(ref db) = db {
        // End the session record
        let _ = db::sessions::end_session(db.pool(), session_id).await;

        // Mark machine as offline
        if let Some(mid) = machine_id {
            let _ = db::machines::mark_machine_offline(db.pool(), mid).await;
        }

        // Log session ended event
        let _ = db::events::log_event(
            db.pool(),
            session_id,
            db::events::EventTypes::SESSION_ENDED,
            None, None, None, None,
        ).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;

            // Database: mark code as completed
            if let Some(ref db) = db {
                let _ = db::support_codes::mark_code_completed(db.pool(), code).await;
            }

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

    info!("Session {} ended", 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 uuid::Uuid;

use crate::proto;
use crate::session::SessionManager;
use crate::db::{self, Database};
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,
    #[serde(default = "default_viewer_name")]
    viewer_name: String,
}

fn default_viewer_name() -> String {
    "Technician".to_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();
    let db = state.db.clone();

    ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, db, 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 viewer_name = params.viewer_name;
    let sessions = state.sessions.clone();
    let db = state.db.clone();

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

/// Handle an agent WebSocket connection
async fn handle_agent_connection(
    socket: WebSocket,
    sessions: SessionManager,
    support_codes: crate::support_codes::SupportCodeManager,
    db: Option<Database>,
    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
    // Persistent agents (no support code) keep their session when disconnected
    let is_persistent = support_code.is_none();
    let (session_id, frame_tx, mut input_rx) = sessions.register_agent(agent_id.clone(), agent_name.clone(), is_persistent).await;

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

    // Database: upsert machine and create session record
    let machine_id = if let Some(ref db) = db {
        match db::machines::upsert_machine(db.pool(), &agent_id, &agent_name, is_persistent).await {
            Ok(machine) => {
                // Create session record
                let _ = db::sessions::create_session(
                    db.pool(),
                    session_id,
                    machine.id,
                    support_code.is_some(),
                    support_code.as_deref(),
                ).await;

                // Log session started event
                let _ = db::events::log_event(
                    db.pool(),
                    session_id,
                    db::events::EventTypes::SESSION_STARTED,
                    None, None, None, None,
                ).await;

                Some(machine.id)
            }
            Err(e) => {
                warn!("Failed to upsert machine in database: {}", e);
                None
            }
        }
    } else {
        None
    };

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

        // Database: update support code
        if let Some(ref db) = db {
            let _ = db::support_codes::mark_code_connected(
                db.pool(),
                code,
                Some(session_id),
                Some(&agent_name),
                Some(&agent_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();
    // Mark agent as disconnected (persistent agents stay in list as offline)
    sessions_cleanup.mark_agent_disconnected(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,
    viewer_name: 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(), viewer_name.clone()).await {
        Some(channels) => channels,
        None => {
            warn!("Session not found: {}", session_id);
            return;
        }
    };

    info!("Viewer {} ({}) joined session: {}", viewer_name, 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);
}

+    // Database: end session and mark machine offline

+        let _ = db::sessions::end_session(db.pool(), session_id).await;

+        if let Some(mid) = machine_id {

+            let _ = db::machines::mark_machine_offline(db.pool(), mid).await;

+            db::events::EventTypes::SESSION_ENDED,

+            // Database: mark code as completed

+                let _ = db::support_codes::mark_code_completed(db.pool(), code).await;