Good, now I need to update the handle_viewer_connection function to:
1. Accept the `db: Option<Database>` parameter
2. Add viewer_joined event logging when viewer joins
3. Add viewer_left event logging when viewer leaves

Also need to update the viewer_ws_handler to pass the db to handle_viewer_connection.


Now let me update the viewer_ws_handler to pass the db parameter:


Edit

/// 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 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 a viewer WebSocket connection
async fn handle_viewer_connection(
    socket: WebSocket,
    sessions: SessionManager,
    db: Option<Database>,
    session_id_str: String,
    viewer_name: 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:
    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 a viewer WebSocket connection
    77→async fn handle_viewer_connection(
    78→    socket: WebSocket,
    79→    sessions: SessionManager,
    80→    db: Option<Database>,
    81→    session_id_str: String,
    82→    viewer_name: String,
    83→) {
    84→
    85→/// Handle an agent WebSocket connection
    86→async fn handle_agent_connection(
    87→    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))
}

/// Handle a viewer WebSocket connection
async fn handle_viewer_connection(
    socket: WebSocket,
    sessions: SessionManager,
    db: Option<Database>,
    session_id_str: String,
    viewer_name: 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,
    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;

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

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

+/// Handle a viewer WebSocket connection


I see there's a problem - I added a new function signature but the old handle_viewer_connection definition is still there. I need to remove the duplicate old definition. Let me look at what happened.


I see the issue - I created a duplicate. Let me fix by removing the old handle_viewer_connection signature that's still there:


   315→
   316→            // Database: mark code as completed
   317→            if let Some(ref db) = db {
   318→                let _ = db::support_codes::mark_code_completed(db.pool(), code).await;
   319→            }
   320→
   321→            info!("Support code {} marked as completed", code);
   322→        }
   323→    }
   324→
   325→    info!("Session {} ended", session_id);
   326→}
   327→
   328→/// Handle a viewer WebSocket connection
   329→async fn handle_viewer_connection(
   330→    socket: WebSocket,
   331→    sessions: SessionManager,
   332→    session_id_str: String,
   333→    viewer_name: String,
   334→) {
   335→    // Parse session ID
   336→    let session_id = match uuid::Uuid::parse_str(&session_id_str) {
   337→        Ok(id) => id,
   338→        Err(_) => {
   339→            warn!("Invalid session ID: {}", session_id_str);
   340→            return;
   341→        }
   342→    };
   343→
   344→    // Generate unique viewer ID

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

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

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


I see the old definition is still there at line 328. I also have an issue where the new signature was inserted in the wrong place (before the agent handler). Let me look at the full state of the file to understand the mess:

1. Line 76-83: New handle_viewer_connection signature I added (but incomplete - just the signature)
2. Line 85+: handle_agent_connection
3. Line 328+: Old handle_viewer_connection

I need to:
1. Remove the incomplete new signature I added (lines 76-83)
2. Update the old handle_viewer_connection (starting at line 328) to include the db parameter and add logging

Let me first remove the partial definition I accidentally added.


I see the issue - I accidentally inserted a partial definition. Let me fix this properly:


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

/// Handle a viewer WebSocket connection
async fn handle_viewer_connection(
    socket: WebSocket,
    sessions: SessionManager,
    db: Option<Database>,
    session_id_str: String,
    viewer_name: String,
) {

/// Handle an agent WebSocket connection

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

/// Handle an agent WebSocket connection


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:
    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,
    79→    sessions: SessionManager,
    80→    support_codes: crate::support_codes::SupportCodeManager,

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

/// Handle an agent WebSocket connection

//! 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 a viewer WebSocket connection
async fn handle_viewer_connection(
    socket: WebSocket,
    sessions: SessionManager,
    db: Option<Database>,
    session_id_str: String,
    viewer_name: 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>,
) {
    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;

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

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

-/// Handle a viewer WebSocket connection


Good, removed the partial definition. Now I need to update the old handle_viewer_connection (now around line 319 after changes) to add the db parameter and logging.


Good. Now let me update the actual handle_viewer_connection function to add the db parameter and logging:


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

/// Handle a viewer WebSocket connection
async fn handle_viewer_connection(
    socket: WebSocket,
    sessions: SessionManager,
    db: Option<Database>,
    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);

    // Database: log viewer joined event
    if let Some(ref db) = db {
        let _ = db::events::log_event(
            db.pool(),
            session_id,
            db::events::EventTypes::VIEWER_JOINED,
            Some(&viewer_id),
            Some(&viewer_name),
            None, None,
        ).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:
   315→
   316→    info!("Session {} ended", session_id);
   317→}
   318→
   319→/// Handle a viewer WebSocket connection
   320→async fn handle_viewer_connection(
   321→    socket: WebSocket,
   322→    sessions: SessionManager,
   323→    db: Option<Database>,
   324→    session_id_str: String,
   325→    viewer_name: String,
   326→) {
   327→    // Parse session ID
   328→    let session_id = match uuid::Uuid::parse_str(&session_id_str) {
   329→        Ok(id) => id,
   330→        Err(_) => {
   331→            warn!("Invalid session ID: {}", session_id_str);
   332→            return;
   333→        }
   334→    };
   335→
   336→    // Generate unique viewer ID
   337→    let viewer_id = Uuid::new_v4().to_string();
   338→
   339→    // Join the session (this sends StartStream to agent if first viewer)
   340→    let (mut frame_rx, input_tx) = match sessions.join_session(session_id, viewer_id.clone(), viewer_name.clone()).await {
   341→        Some(channels) => channels,
   342→        None => {
   343→            warn!("Session not found: {}", session_id);
   344→            return;
   345→        }
   346→    };
   347→
   348→    info!("Viewer {} ({}) joined session: {}", viewer_name, viewer_id, session_id);
   349→
   350→    // Database: log viewer joined event
   351→    if let Some(ref db) = db {
   352→        let _ = db::events::log_event(
   353→            db.pool(),
   354→            session_id,
   355→            db::events::EventTypes::VIEWER_JOINED,
   356→            Some(&viewer_id),
   357→            Some(&viewer_name),
   358→            None, None,
   359→        ).await;
   360→    }
   361→
   362→    let (mut ws_sender, mut ws_receiver) = socket.split();
   363→
   364→    // Task to forward frames from agent to this viewer

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

/// Handle a viewer WebSocket connection
async fn handle_viewer_connection(
    socket: WebSocket,
    sessions: SessionManager,
    db: Option<Database>,
    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);

    // Database: log viewer joined event
    if let Some(ref db) = db {
        let _ = db::events::log_event(
            db.pool(),
            session_id,
            db::events::EventTypes::VIEWER_JOINED,
            Some(&viewer_id),
            Some(&viewer_name),
            None, None,
        ).await;
    }

async fn handle_viewer_connection(
socket: WebSocket,
sessions: SessionManager,
+    db: Option<Database>,
session_id_str: String,
viewer_name: String,
) {

+    // Database: log viewer joined event

+            db::events::EventTypes::VIEWER_JOINED,


Good progress! Now I need to add the viewer_left logging at the end of the handle_viewer_connection function. Let me find the cleanup section.


Now I need to add the viewer_left logging at the cleanup section:


   385→                            Some(proto::message::Payload::KeyEvent(_)) |
   386→                            Some(proto::message::Payload::SpecialKey(_)) => {
   387→                                // Forward input to agent
   388→                                let _ = input_tx.send(data.to_vec()).await;
   389→                            }
   390→                            Some(proto::message::Payload::ChatMessage(chat)) => {
   391→                                // Forward chat message to agent
   392→                                info!("Chat from technician: {}", chat.content);
   393→                                let _ = input_tx.send(data.to_vec()).await;
   394→                            }
   395→                            _ => {}
   396→                        }
   397→                    }
   398→                    Err(e) => {
   399→                        warn!("Failed to decode viewer message: {}", e);
   400→                    }
   401→                }
   402→            }
   403→            Ok(Message::Close(_)) => {
   404→                info!("Viewer {} disconnected from session: {}", viewer_id, session_id);
   405→                break;
   406→            }
   407→            Ok(_) => {}
   408→            Err(e) => {
   409→                error!("WebSocket error from viewer {}: {}", viewer_id, e);
   410→                break;
   411→            }
   412→        }
   413→    }
   414→
   415→    // Cleanup (this sends StopStream to agent if last viewer)
   416→    frame_forward.abort();
   417→    sessions_cleanup.leave_session(session_id, &viewer_id_cleanup).await;
   418→    info!("Viewer {} left session: {}", viewer_id_cleanup, session_id);
   419→}
   420→

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

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


Now I need to update the cleanup section at the end of handle_viewer_connection to add viewer_left logging. I also notice there's a `sessions_cleanup` variable being used - I need to make sure the db is also cloned for cleanup. Let me find where those clones happen.


   362→    let (mut ws_sender, mut ws_receiver) = socket.split();
   363→
   364→    // Task to forward frames from agent to this viewer
   365→    let frame_forward = tokio::spawn(async move {
   366→        while let Ok(frame_data) = frame_rx.recv().await {
   367→            if ws_sender.send(Message::Binary(frame_data.into())).await.is_err() {
   368→                break;
   369→            }
   370→        }
   371→    });
   372→
   373→    let sessions_cleanup = sessions.clone();
   374→    let viewer_id_cleanup = viewer_id.clone();
   375→
   376→    // Main loop: receive input from viewer and forward to agent
   377→    while let Some(msg) = ws_receiver.next().await {
   378→        match msg {
   379→            Ok(Message::Binary(data)) => {
   380→                // Try to decode as protobuf message
   381→                match proto::Message::decode(data.as_ref()) {
   382→                    Ok(proto_msg) => {
   383→                        match &proto_msg.payload {
   384→                            Some(proto::message::Payload::MouseEvent(_)) |
   385→                            Some(proto::message::Payload::KeyEvent(_)) |
   386→                            Some(proto::message::Payload::SpecialKey(_)) => {

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


Now I need to add the viewer_name_cleanup clone and update the cleanup section to log the viewer_left event.


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

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

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

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


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:
   369→            }
   370→        }
   371→    });
   372→
   373→    let sessions_cleanup = sessions.clone();
   374→    let viewer_id_cleanup = viewer_id.clone();
   375→    let viewer_name_cleanup = viewer_name.clone();
   376→
   377→    // Main loop: receive input from viewer and forward to agent
   378→    while let Some(msg) = ws_receiver.next().await {
   379→        match msg {
   380→            Ok(Message::Binary(data)) => {
   381→                // Try to decode as protobuf message

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

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

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

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

/// Handle a viewer WebSocket connection
async fn handle_viewer_connection(
    socket: WebSocket,
    sessions: SessionManager,
    db: Option<Database>,
    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);

    // Database: log viewer joined event
    if let Some(ref db) = db {
        let _ = db::events::log_event(
            db.pool(),
            session_id,
            db::events::EventTypes::VIEWER_JOINED,
            Some(&viewer_id),
            Some(&viewer_name),
            None, None,
        ).await;
    }

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

+    let viewer_name_cleanup = viewer_name.clone();


Good, now I need to update the cleanup section at the end to add the viewer_left database logging.


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

    // Cleanup (this sends StopStream to agent if last viewer)
    frame_forward.abort();
    sessions_cleanup.leave_session(session_id, &viewer_id_cleanup).await;

    // Database: log viewer left event
    if let Some(ref db) = db {
        let _ = db::events::log_event(
            db.pool(),
            session_id,
            db::events::EventTypes::VIEWER_LEFT,
            Some(&viewer_id_cleanup),
            Some(&viewer_name_cleanup),
            None, None,
        ).await;
    }

    info!("Viewer {} left session: {}", viewer_id_cleanup, 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:
   412→            }
   413→        }
   414→    }
   415→
   416→    // Cleanup (this sends StopStream to agent if last viewer)
   417→    frame_forward.abort();
   418→    sessions_cleanup.leave_session(session_id, &viewer_id_cleanup).await;
   419→
   420→    // Database: log viewer left event
   421→    if let Some(ref db) = db {
   422→        let _ = db::events::log_event(
   423→            db.pool(),
   424→            session_id,
   425→            db::events::EventTypes::VIEWER_LEFT,
   426→            Some(&viewer_id_cleanup),
   427→            Some(&viewer_name_cleanup),
   428→            None, None,
   429→        ).await;
   430→    }
   431→
   432→    info!("Viewer {} left session: {}", viewer_id_cleanup, session_id);
   433→}
   434→

    // Cleanup (this sends StopStream to agent if last viewer)
    frame_forward.abort();
    sessions_cleanup.leave_session(session_id, &viewer_id_cleanup).await;

    // Database: log viewer left event
    if let Some(ref db) = db {
        let _ = db::events::log_event(
            db.pool(),
            session_id,
            db::events::EventTypes::VIEWER_LEFT,
            Some(&viewer_id_cleanup),
            Some(&viewer_name_cleanup),
            None, None,
        ).await;
    }

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

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

use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        Query, State,
    },
    response::IntoResponse,
};
use futures_util::{SinkExt, StreamExt};
use prost::Message as ProstMessage;
use serde::Deserialize;
use tracing::{error, info, warn};
use 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;

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

/// Handle a viewer WebSocket connection
async fn handle_viewer_connection(
    socket: WebSocket,
    sessions: SessionManager,
    db: Option<Database>,
    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);

    // Database: log viewer joined event
    if let Some(ref db) = db {
        let _ = db::events::log_event(
            db.pool(),
            session_id,
            db::events::EventTypes::VIEWER_JOINED,
            Some(&viewer_id),
            Some(&viewer_name),
            None, None,
        ).await;
    }

    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();
    let viewer_name_cleanup = viewer_name.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);
}

+            db::events::EventTypes::VIEWER_LEFT,

+            Some(&viewer_name_cleanup),


The relay/mod.rs updates are complete. Now let me update the todo list and try to build to make sure everything compiles.


Now let me update the todo list and try building to verify everything compiles:


TodoWrite

Create guruconnect database on build server

Create migration file (001_initial_schema.sql)

Expand db module with connection pool

Create db/machines.rs

Create db/sessions.rs

Create db/events.rs

Create db/support_codes.rs

Update main.rs with DB init

Update session/mod.rs with restore method

Update relay/mod.rs with event logging

Build and verify compilation

Deploy and test


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

Create guruconnect database on build server

Create migration file (001_initial_schema.sql)

Expand db module with connection pool

Create db/machines.rs

Create db/sessions.rs

Create db/events.rs

Create db/support_codes.rs

Update main.rs with DB init

Update session/mod.rs with restore method

Update relay/mod.rs with event logging

Deploy and test

Expand db module with connection pool

Create db/machines.rs

Create db/sessions.rs

Create db/events.rs

Create db/support_codes.rs

Update main.rs with DB init

Update relay/mod.rs with event logging

Build and verify compilation

Deploy and test


Let me try building the server to see if there are any compilation errors.


Bash

cd /c/Users/MikeSwanson/Claude/guru-connect && cargo build --release -p guruconnect-server 2>&1 | head -100

Build GuruConnect server


warning: profiles for the non root package will be ignored, specify profiles at the workspace root:
package:   C:\Users\MikeSwanson\Claude\guru-connect\agent\Cargo.toml
workspace: C:\Users\MikeSwanson\Claude\guru-connect\Cargo.toml
warning: profiles for the non root package will be ignored, specify profiles at the workspace root:
package:   C:\Users\MikeSwanson\Claude\guru-connect\server\Cargo.toml
workspace: C:\Users\MikeSwanson\Claude\guru-connect\Cargo.toml
   Compiling guruconnect-server v0.1.0 (C:\Users\MikeSwanson\Claude\guru-connect\server)
warning: unused import: `machines::*`
  --> server\src\db\mod.rs:16:9
   |
16 | pub use machines::*;
   |         ^^^^^^^^^^^
   |
   = note: `#[warn(unused_imports)]` (part of `#[warn(unused)]`) on by default

warning: unused import: `sessions::*`
  --> server\src\db\mod.rs:17:9
   |
17 | pub use sessions::*;
   |         ^^^^^^^^^^^

warning: unused import: `events::*`
  --> server\src\db\mod.rs:18:9
   |
18 | pub use events::*;
   |         ^^^^^^^^^

warning: unused import: `support_codes::*`
  --> server\src\db\mod.rs:19:9
   |
19 | pub use support_codes::*;
   |         ^^^^^^^^^^^^^^^^

error[E0061]: this function takes 3 arguments but 2 arguments were supplied
   --> server\src\relay\mod.rs:286:17
    |
286 |         let _ = db::sessions::end_session(db.pool(), session_id).await;
    |                 ^^^^^^^^^^^^^^^^^^^^^^^^^----------------------- argument #3 of type `&str` is missing
    |
note: function defined here
   --> server\src\db\sessions.rs:45:14
    |
 45 | pub async fn end_session(
    |              ^^^^^^^^^^^
...
 48 |     status: &str,  // 'ended' or 'disconnected' or 'timeout'
    |     ------------
help: provide the argument
    |
286 |         let _ = db::sessions::end_session(db.pool(), session_id, /* &str */).await;
    |                                                                ++++++++++++

error[E0308]: mismatched types
   --> server\src\relay\mod.rs:290:67
    |
290 |             let _ = db::machines::mark_machine_offline(db.pool(), mid).await;
    |                     ----------------------------------            ^^^ expected `&str`, found `Uuid`
    |                     |
    |                     arguments to this function are incorrect
    |
note: function defined here
   --> server\src\db\machines.rs:103:14
    |
103 | pub async fn mark_machine_offline(pool: &PgPool, agent_id: &str) -> Result<(), sqlx::Error> {
    |              ^^^^^^^^^^^^^^^^^^^^                --------------

Some errors have detailed explanations: E0061, E0308.
For more information about an error, try `rustc --explain E0061`.
warning: `guruconnect-server` (bin "guruconnect-server") generated 4 warnings
error: could not compile `guruconnect-server` (bin "guruconnect-server") due to 2 previous errors; 4 warnings emitted


There are two errors:
1. end_session needs a third argument (status string)
2. mark_machine_offline expects agent_id string, not machine UUID

Let me fix these issues.


I need to fix two issues in the relay/mod.rs - the function signatures don't match:


Edit

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

    // 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, "ended").await;

        // Mark machine as offline
        let _ = db::machines::mark_machine_offline(db.pool(), &agent_id).await;