     1→//! Server configuration
     2→
     3→use anyhow::Result;
     4→use serde::Deserialize;
     5→use std::env;
     6→
     7→#[derive(Debug, Clone, Deserialize)]
     8→pub struct Config {
     9→    /// Address to listen on (e.g., "0.0.0.0:8080")
    10→    pub listen_addr: String,
    11→
    12→    /// Database URL (optional for MVP)
    13→    pub database_url: Option<String>,
    14→
    15→    /// JWT secret for authentication
    16→    pub jwt_secret: Option<String>,
    17→
    18→    /// Enable debug logging
    19→    pub debug: bool,
    20→}
    21→
    22→impl Config {
    23→    /// Load configuration from environment variables
    24→    pub fn load() -> Result<Self> {
    25→        Ok(Self {
    26→            listen_addr: env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:8080".to_string()),
    27→            database_url: env::var("DATABASE_URL").ok(),
    28→            jwt_secret: env::var("JWT_SECRET").ok(),
    29→            debug: env::var("DEBUG")
    30→                .map(|v| v == "1" || v.to_lowercase() == "true")
    31→                .unwrap_or(false),
    32→        })
    33→    }
    34→}
    35→
    36→impl Default for Config {
    37→    fn default() -> Self {
    38→        Self {
    39→            listen_addr: "0.0.0.0:8080".to_string(),
    40→            database_url: None,
    41→            jwt_secret: None,
    42→            debug: false,
    43→        }
    44→    }
    45→}
    46→

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

//! Server configuration

use anyhow::Result;
use serde::Deserialize;
use std::env;

#[derive(Debug, Clone, Deserialize)]
pub struct Config {
    /// Address to listen on (e.g., "0.0.0.0:8080")
    pub listen_addr: String,

    /// Database URL (optional for MVP)
    pub database_url: Option<String>,

    /// JWT secret for authentication
    pub jwt_secret: Option<String>,

    /// Enable debug logging
    pub debug: bool,
}

impl Config {
    /// Load configuration from environment variables
    pub fn load() -> Result<Self> {
        Ok(Self {
            listen_addr: env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:8080".to_string()),
            database_url: env::var("DATABASE_URL").ok(),
            jwt_secret: env::var("JWT_SECRET").ok(),
            debug: env::var("DEBUG")
                .map(|v| v == "1" || v.to_lowercase() == "true")
                .unwrap_or(false),
        })
    }
}

impl Default for Config {
    fn default() -> Self {
        Self {
            listen_addr: "0.0.0.0:8080".to_string(),
            database_url: None,
            jwt_secret: None,
            debug: false,
        }
    }
}


Good. Now let me update the config to add database_max_connections, then update main.rs with database initialization.


Edit

#[derive(Debug, Clone, Deserialize)]
pub struct Config {
    /// Address to listen on (e.g., "0.0.0.0:8080")
    pub listen_addr: String,

    /// Database URL (optional for MVP)
    pub database_url: Option<String>,

    /// JWT secret for authentication
    pub jwt_secret: Option<String>,

    /// Enable debug logging
    pub debug: bool,
}

#[derive(Debug, Clone, Deserialize)]
pub struct Config {
    /// Address to listen on (e.g., "0.0.0.0:8080")
    pub listen_addr: String,

    /// Database URL (optional - server works without it)
    pub database_url: Option<String>,

    /// Maximum database connections in pool
    pub database_max_connections: u32,

    /// JWT secret for authentication
    pub jwt_secret: Option<String>,

    /// Enable debug logging
    pub debug: bool,
}


The file C:\Users\MikeSwanson\Claude\guru-connect\server\src\config.rs has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
     3→use anyhow::Result;
     4→use serde::Deserialize;
     5→use std::env;
     6→
     7→#[derive(Debug, Clone, Deserialize)]
     8→pub struct Config {
     9→    /// Address to listen on (e.g., "0.0.0.0:8080")
    10→    pub listen_addr: String,
    11→
    12→    /// Database URL (optional - server works without it)
    13→    pub database_url: Option<String>,
    14→
    15→    /// Maximum database connections in pool
    16→    pub database_max_connections: u32,
    17→
    18→    /// JWT secret for authentication
    19→    pub jwt_secret: Option<String>,
    20→
    21→    /// Enable debug logging
    22→    pub debug: bool,
    23→}
    24→
    25→impl Config {
    26→    /// Load configuration from environment variables
    27→    pub fn load() -> Result<Self> {

#[derive(Debug, Clone, Deserialize)]
pub struct Config {
    /// Address to listen on (e.g., "0.0.0.0:8080")
    pub listen_addr: String,

    /// Database URL (optional - server works without it)
    pub database_url: Option<String>,

    /// Maximum database connections in pool
    pub database_max_connections: u32,

    /// JWT secret for authentication
    pub jwt_secret: Option<String>,

    /// Enable debug logging
    pub debug: bool,
}

//! Server configuration

use anyhow::Result;
use serde::Deserialize;
use std::env;

#[derive(Debug, Clone, Deserialize)]
pub struct Config {
    /// Address to listen on (e.g., "0.0.0.0:8080")
    pub listen_addr: String,

    /// Database URL (optional for MVP)
    pub database_url: Option<String>,

    /// JWT secret for authentication
    pub jwt_secret: Option<String>,

    /// Enable debug logging
    pub debug: bool,
}

impl Config {
    /// Load configuration from environment variables
    pub fn load() -> Result<Self> {
        Ok(Self {
            listen_addr: env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:8080".to_string()),
            database_url: env::var("DATABASE_URL").ok(),
            jwt_secret: env::var("JWT_SECRET").ok(),
            debug: env::var("DEBUG")
                .map(|v| v == "1" || v.to_lowercase() == "true")
                .unwrap_or(false),
        })
    }
}

impl Default for Config {
    fn default() -> Self {
        Self {
            listen_addr: "0.0.0.0:8080".to_string(),
            database_url: None,
            jwt_secret: None,
            debug: false,
        }
    }
}

-    /// Database URL (optional for MVP)

+    /// Database URL (optional - server works without it)

+    /// Maximum database connections in pool


Good. Now let me update the load() and default() methods to include database_max_connections:


impl Config {
    /// Load configuration from environment variables
    pub fn load() -> Result<Self> {
        Ok(Self {
            listen_addr: env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:8080".to_string()),
            database_url: env::var("DATABASE_URL").ok(),
            jwt_secret: env::var("JWT_SECRET").ok(),
            debug: env::var("DEBUG")
                .map(|v| v == "1" || v.to_lowercase() == "true")
                .unwrap_or(false),
        })
    }
}

impl Default for Config {
    fn default() -> Self {
        Self {
            listen_addr: "0.0.0.0:8080".to_string(),
            database_url: None,
            jwt_secret: None,
            debug: false,
        }
    }
}

impl Config {
    /// Load configuration from environment variables
    pub fn load() -> Result<Self> {
        Ok(Self {
            listen_addr: env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:8080".to_string()),
            database_url: env::var("DATABASE_URL").ok(),
            database_max_connections: env::var("DATABASE_MAX_CONNECTIONS")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(5),
            jwt_secret: env::var("JWT_SECRET").ok(),
            debug: env::var("DEBUG")
                .map(|v| v == "1" || v.to_lowercase() == "true")
                .unwrap_or(false),
        })
    }
}

impl Default for Config {
    fn default() -> Self {
        Self {
            listen_addr: "0.0.0.0:8080".to_string(),
            database_url: None,
            database_max_connections: 5,
            jwt_secret: None,
            debug: false,
        }
    }
}


The file C:\Users\MikeSwanson\Claude\guru-connect\server\src\config.rs has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
    21→    /// Enable debug logging
    22→    pub debug: bool,
    23→}
    24→
    25→impl Config {
    26→    /// Load configuration from environment variables
    27→    pub fn load() -> Result<Self> {
    28→        Ok(Self {
    29→            listen_addr: env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:8080".to_string()),
    30→            database_url: env::var("DATABASE_URL").ok(),
    31→            database_max_connections: env::var("DATABASE_MAX_CONNECTIONS")
    32→                .ok()
    33→                .and_then(|v| v.parse().ok())
    34→                .unwrap_or(5),
    35→            jwt_secret: env::var("JWT_SECRET").ok(),
    36→            debug: env::var("DEBUG")
    37→                .map(|v| v == "1" || v.to_lowercase() == "true")
    38→                .unwrap_or(false),
    39→        })
    40→    }
    41→}
    42→
    43→impl Default for Config {
    44→    fn default() -> Self {
    45→        Self {
    46→            listen_addr: "0.0.0.0:8080".to_string(),
    47→            database_url: None,
    48→            database_max_connections: 5,
    49→            jwt_secret: None,
    50→            debug: false,
    51→        }
    52→    }
    53→}
    54→

impl Config {
    /// Load configuration from environment variables
    pub fn load() -> Result<Self> {
        Ok(Self {
            listen_addr: env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:8080".to_string()),
            database_url: env::var("DATABASE_URL").ok(),
            database_max_connections: env::var("DATABASE_MAX_CONNECTIONS")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(5),
            jwt_secret: env::var("JWT_SECRET").ok(),
            debug: env::var("DEBUG")
                .map(|v| v == "1" || v.to_lowercase() == "true")
                .unwrap_or(false),
        })
    }
}

impl Default for Config {
    fn default() -> Self {
        Self {
            listen_addr: "0.0.0.0:8080".to_string(),
            database_url: None,
            database_max_connections: 5,
            jwt_secret: None,
            debug: false,
        }
    }
}

//! Server configuration

use anyhow::Result;
use serde::Deserialize;
use std::env;

#[derive(Debug, Clone, Deserialize)]
pub struct Config {
    /// Address to listen on (e.g., "0.0.0.0:8080")
    pub listen_addr: String,

    /// Database URL (optional - server works without it)
    pub database_url: Option<String>,

    /// Maximum database connections in pool
    pub database_max_connections: u32,

    /// JWT secret for authentication
    pub jwt_secret: Option<String>,

    /// Enable debug logging
    pub debug: bool,
}

impl Config {
    /// Load configuration from environment variables
    pub fn load() -> Result<Self> {
        Ok(Self {
            listen_addr: env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:8080".to_string()),
            database_url: env::var("DATABASE_URL").ok(),
            jwt_secret: env::var("JWT_SECRET").ok(),
            debug: env::var("DEBUG")
                .map(|v| v == "1" || v.to_lowercase() == "true")
                .unwrap_or(false),
        })
    }
}

impl Default for Config {
    fn default() -> Self {
        Self {
            listen_addr: "0.0.0.0:8080".to_string(),
            database_url: None,
            jwt_secret: None,
            debug: false,
        }
    }
}

+            database_max_connections: env::var("DATABASE_MAX_CONNECTIONS")

+                .and_then(|v| v.parse().ok())

+            database_max_connections: 5,


Good. Now let me update main.rs to add database initialization and pass it to AppState.


Now let me update main.rs:


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

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


The file C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
    32→use serde::Deserialize;
    33→
    34→use support_codes::{SupportCodeManager, CreateCodeRequest, SupportCode, CodeValidation};
    35→
    36→/// Application state
    37→#[derive(Clone)]
    38→pub struct AppState {
    39→    sessions: session::SessionManager,
    40→    support_codes: SupportCodeManager,
    41→    db: Option<db::Database>,
    42→}
    43→
    44→#[tokio::main]
    45→async fn main() -> Result<()> {
    46→    // Initialize logging

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        // State
        .with_state(state)

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

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

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

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

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

    Ok(())
}

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

// Support code API handlers

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

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

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

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

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

// Session API handlers (updated to use AppState)

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

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

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

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

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

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

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

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

pub struct AppState {
sessions: session::SessionManager,
support_codes: SupportCodeManager,
+    db: Option<db::Database>,
}
#[tokio::main]


Good. Now I need to update the main function to initialize the database. Let me replace the create application state section.


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

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

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

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

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

    // Initialize database if configured
    let database = if let Some(ref db_url) = config.database_url {
        match db::Database::connect(db_url, config.database_max_connections).await {
            Ok(db) => {
                // Run migrations
                if let Err(e) = db.migrate().await {
                    tracing::error!("Failed to run migrations: {}", e);
                    return Err(e);
                }
                Some(db)
            }
            Err(e) => {
                tracing::warn!("Failed to connect to database: {}. Running without persistence.", e);
                None
            }
        }
    } else {
        info!("No DATABASE_URL set, running without persistence");
        None
    };

    // Create session manager
    let sessions = session::SessionManager::new();

    // Restore persistent machines from database
    if let Some(ref db) = database {
        match db::machines::get_all_machines(db.pool()).await {
            Ok(machines) => {
                info!("Restoring {} persistent machines from database", machines.len());
                for machine in machines {
                    sessions.restore_offline_machine(&machine.agent_id, &machine.hostname).await;
                }
            }
            Err(e) => {
                tracing::warn!("Failed to restore machines: {}", e);
            }
        }
    }

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


The file C:\Users\MikeSwanson\Claude\guru-connect\server\src\main.rs has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
    50→        .init();
    51→
    52→    info!("GuruConnect Server v{}", env!("CARGO_PKG_VERSION"));
    53→
    54→    // Load configuration
    55→    let config = config::Config::load()?;
    56→
    57→    // Use port 3002 for GuruConnect
    58→    let listen_addr = std::env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:3002".to_string());
    59→    info!("Loaded configuration, listening on {}", listen_addr);
    60→
    61→    // Initialize database if configured
    62→    let database = if let Some(ref db_url) = config.database_url {
    63→        match db::Database::connect(db_url, config.database_max_connections).await {
    64→            Ok(db) => {
    65→                // Run migrations
    66→                if let Err(e) = db.migrate().await {
    67→                    tracing::error!("Failed to run migrations: {}", e);
    68→                    return Err(e);
    69→                }
    70→                Some(db)
    71→            }
    72→            Err(e) => {
    73→                tracing::warn!("Failed to connect to database: {}. Running without persistence.", e);
    74→                None
    75→            }
    76→        }
    77→    } else {
    78→        info!("No DATABASE_URL set, running without persistence");
    79→        None
    80→    };
    81→
    82→    // Create session manager
    83→    let sessions = session::SessionManager::new();
    84→
    85→    // Restore persistent machines from database
    86→    if let Some(ref db) = database {
    87→        match db::machines::get_all_machines(db.pool()).await {
    88→            Ok(machines) => {
    89→                info!("Restoring {} persistent machines from database", machines.len());
    90→                for machine in machines {
    91→                    sessions.restore_offline_machine(&machine.agent_id, &machine.hostname).await;
    92→                }
    93→            }
    94→            Err(e) => {
    95→                tracing::warn!("Failed to restore machines: {}", e);
    96→            }
    97→        }
    98→    }
    99→
   100→    // Create application state
   101→    let state = AppState {
   102→        sessions,
   103→        support_codes: SupportCodeManager::new(),
   104→        db: database,
   105→    };
   106→
   107→    // Build router
   108→    let app = Router::new()
   109→        // Health check

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

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

    // Initialize database if configured
    let database = if let Some(ref db_url) = config.database_url {
        match db::Database::connect(db_url, config.database_max_connections).await {
            Ok(db) => {
                // Run migrations
                if let Err(e) = db.migrate().await {
                    tracing::error!("Failed to run migrations: {}", e);
                    return Err(e);
                }
                Some(db)
            }
            Err(e) => {
                tracing::warn!("Failed to connect to database: {}. Running without persistence.", e);
                None
            }
        }
    } else {
        info!("No DATABASE_URL set, running without persistence");
        None
    };

    // Create session manager
    let sessions = session::SessionManager::new();

    // Restore persistent machines from database
    if let Some(ref db) = database {
        match db::machines::get_all_machines(db.pool()).await {
            Ok(machines) => {
                info!("Restoring {} persistent machines from database", machines.len());
                for machine in machines {
                    sessions.restore_offline_machine(&machine.agent_id, &machine.hostname).await;
                }
            }
            Err(e) => {
                tracing::warn!("Failed to restore machines: {}", e);
            }
        }
    }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        // State
        .with_state(state)

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

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

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

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

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

    Ok(())
}

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

// Support code API handlers

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

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

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

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

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

// Session API handlers (updated to use AppState)

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

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

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

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

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

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

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

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

+    // Initialize database if configured

+    let database = if let Some(ref db_url) = config.database_url {

+        match db::Database::connect(db_url, config.database_max_connections).await {

+                if let Err(e) = db.migrate().await {

+                    tracing::error!("Failed to run migrations: {}", e);

+                tracing::warn!("Failed to connect to database: {}. Running without persistence.", e);

+        info!("No DATABASE_URL set, running without persistence");

+    let sessions = session::SessionManager::new();

+    // Restore persistent machines from database

+        match db::machines::get_all_machines(db.pool()).await {

+                info!("Restoring {} persistent machines from database", machines.len());

+                for machine in machines {

+                    sessions.restore_offline_machine(&machine.agent_id, &machine.hostname).await;

+                tracing::warn!("Failed to restore machines: {}", e);

-        sessions: session::SessionManager::new(),


Good. Now I need to add the restore_offline_machine method to the SessionManager. Let me read the session/mod.rs file and add it.


Now let me add the `restore_offline_machine` method to SessionManager:


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

Deploy and test


Edit

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

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

impl SessionManager {
    /// Restore a machine as an offline session (called on startup from database)
    pub async fn restore_offline_machine(&self, agent_id: &str, hostname: &str) -> SessionId {
        let session_id = Uuid::new_v4();
        let now = chrono::Utc::now();

        let session = Session {
            id: session_id,
            agent_id: agent_id.to_string(),
            agent_name: hostname.to_string(),
            started_at: now,
            viewer_count: 0,
            viewers: Vec::new(),
            is_streaming: false,
            is_online: false,  // Offline until agent reconnects
            is_persistent: true,
            last_heartbeat: now,
            os_version: None,
            is_elevated: false,
            uptime_secs: 0,
            display_count: 1,
        };

        // Create placeholder channels (will be replaced on reconnect)
        let (frame_tx, _) = broadcast::channel(16);
        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64);

        let session_data = SessionData {
            info: session,
            frame_tx,
            input_tx,
            input_rx: Some(input_rx),
            viewers: HashMap::new(),
            last_heartbeat_instant: Instant::now(),
        };

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

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

        tracing::info!("Restored offline machine: {} ({})", hostname, agent_id);
        session_id
    }
}


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

Deploy and test


The file C:\Users\MikeSwanson\Claude\guru-connect\server\src\session\mod.rs has been updated. Here's the result of running `cat -n` on a snippet of the edited file:
   384→        sessions.values().map(|s| s.info.clone()).collect()
   385→    }
   386→}
   387→
   388→impl Default for SessionManager {
   389→    fn default() -> Self {
   390→        Self::new()
   391→    }
   392→}
   393→
   394→impl SessionManager {
   395→    /// Restore a machine as an offline session (called on startup from database)
   396→    pub async fn restore_offline_machine(&self, agent_id: &str, hostname: &str) -> SessionId {
   397→        let session_id = Uuid::new_v4();
   398→        let now = chrono::Utc::now();
   399→
   400→        let session = Session {
   401→            id: session_id,
   402→            agent_id: agent_id.to_string(),
   403→            agent_name: hostname.to_string(),
   404→            started_at: now,
   405→            viewer_count: 0,
   406→            viewers: Vec::new(),
   407→            is_streaming: false,
   408→            is_online: false,  // Offline until agent reconnects
   409→            is_persistent: true,
   410→            last_heartbeat: now,
   411→            os_version: None,
   412→            is_elevated: false,
   413→            uptime_secs: 0,
   414→            display_count: 1,
   415→        };
   416→
   417→        // Create placeholder channels (will be replaced on reconnect)
   418→        let (frame_tx, _) = broadcast::channel(16);
   419→        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64);
   420→
   421→        let session_data = SessionData {
   422→            info: session,
   423→            frame_tx,
   424→            input_tx,
   425→            input_rx: Some(input_rx),
   426→            viewers: HashMap::new(),
   427→            last_heartbeat_instant: Instant::now(),
   428→        };
   429→
   430→        let mut sessions = self.sessions.write().await;
   431→        sessions.insert(session_id, session_data);
   432→
   433→        let mut agents = self.agents.write().await;
   434→        agents.insert(agent_id.to_string(), session_id);
   435→
   436→        tracing::info!("Restored offline machine: {} ({})", hostname, agent_id);
   437→        session_id
   438→    }
   439→}
   440→

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

impl SessionManager {
    /// Restore a machine as an offline session (called on startup from database)
    pub async fn restore_offline_machine(&self, agent_id: &str, hostname: &str) -> SessionId {
        let session_id = Uuid::new_v4();
        let now = chrono::Utc::now();

        let session = Session {
            id: session_id,
            agent_id: agent_id.to_string(),
            agent_name: hostname.to_string(),
            started_at: now,
            viewer_count: 0,
            viewers: Vec::new(),
            is_streaming: false,
            is_online: false,  // Offline until agent reconnects
            is_persistent: true,
            last_heartbeat: now,
            os_version: None,
            is_elevated: false,
            uptime_secs: 0,
            display_count: 1,
        };

        // Create placeholder channels (will be replaced on reconnect)
        let (frame_tx, _) = broadcast::channel(16);
        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64);

        let session_data = SessionData {
            info: session,
            frame_tx,
            input_tx,
            input_rx: Some(input_rx),
            viewers: HashMap::new(),
            last_heartbeat_instant: Instant::now(),
        };

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

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

        tracing::info!("Restored offline machine: {} ({})", hostname, agent_id);
        session_id
    }
}

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

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

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

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

/// Unique identifier for a viewer
pub type ViewerId = String;

/// Information about a connected viewer/technician
#[derive(Debug, Clone)]
pub struct ViewerInfo {
    pub id: ViewerId,
    pub name: String,
    pub connected_at: chrono::DateTime<chrono::Utc>,
}

/// Heartbeat timeout (90 seconds - 3x the agent's 30 second interval)
const HEARTBEAT_TIMEOUT_SECS: u64 = 90;

/// Session state
#[derive(Debug, Clone)]
pub struct Session {
    pub id: SessionId,
    pub agent_id: AgentId,
    pub agent_name: String,
    pub started_at: chrono::DateTime<chrono::Utc>,
    pub viewer_count: usize,
    pub viewers: Vec<ViewerInfo>,  // List of connected technicians
    pub is_streaming: bool,
    pub is_online: bool,  // Whether agent is currently connected
    pub is_persistent: bool,  // Persistent agent (no support code) vs support session
    pub last_heartbeat: chrono::DateTime<chrono::Utc>,
    // Agent status info
    pub os_version: Option<String>,
    pub is_elevated: bool,
    pub uptime_secs: i64,
    pub display_count: i32,
}

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

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

/// Internal session data with channels
struct SessionData {
    info: Session,
    /// Channel for video frames (agent -> viewers)
    frame_tx: FrameSender,
    /// Channel for input events (viewer -> agent)
    input_tx: InputSender,
    input_rx: Option<InputReceiver>,
    /// Map of connected viewers (id -> info)
    viewers: HashMap<ViewerId, ViewerInfo>,
    /// Instant for heartbeat tracking
    last_heartbeat_instant: Instant,
}

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

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

    /// Register a new agent and create a session
    /// If agent was previously connected (offline session exists), reuse that session
    pub async fn register_agent(&self, agent_id: AgentId, agent_name: String, is_persistent: bool) -> (SessionId, FrameSender, InputReceiver) {
        // Check if this agent already has an offline session (reconnecting)
        {
            let agents = self.agents.read().await;
            if let Some(&existing_session_id) = agents.get(&agent_id) {
                let mut sessions = self.sessions.write().await;
                if let Some(session_data) = sessions.get_mut(&existing_session_id) {
                    if !session_data.info.is_online {
                        // Reuse existing session - mark as online and create new channels
                        tracing::info!("Agent {} reconnecting to existing session {}", agent_id, existing_session_id);

                        let (frame_tx, _) = broadcast::channel(16);
                        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64);

                        session_data.info.is_online = true;
                        session_data.info.last_heartbeat = chrono::Utc::now();
                        session_data.info.agent_name = agent_name; // Update name in case it changed
                        session_data.frame_tx = frame_tx.clone();
                        session_data.input_tx = input_tx;
                        session_data.last_heartbeat_instant = Instant::now();

                        return (existing_session_id, frame_tx, input_rx);
                    }
                }
            }
        }

        // Create new session
        let session_id = Uuid::new_v4();

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

        let now = chrono::Utc::now();
        let session = Session {
            id: session_id,
            agent_id: agent_id.clone(),
            agent_name,
            started_at: now,
            viewer_count: 0,
            viewers: Vec::new(),
            is_streaming: false,
            is_online: true,
            is_persistent,
            last_heartbeat: now,
            os_version: None,
            is_elevated: false,
            uptime_secs: 0,
            display_count: 1,
        };

        let session_data = SessionData {
            info: session,
            frame_tx: frame_tx.clone(),
            input_tx,
            input_rx: None,
            viewers: HashMap::new(),
            last_heartbeat_instant: Instant::now(),
        };

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

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

        (session_id, frame_tx, input_rx)
    }

    /// Update agent status from heartbeat or status message
    pub async fn update_agent_status(
        &self,
        session_id: SessionId,
        os_version: Option<String>,
        is_elevated: bool,
        uptime_secs: i64,
        display_count: i32,
        is_streaming: bool,
    ) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.get_mut(&session_id) {
            session_data.info.last_heartbeat = chrono::Utc::now();
            session_data.last_heartbeat_instant = Instant::now();
            session_data.info.is_streaming = is_streaming;
            if let Some(os) = os_version {
                session_data.info.os_version = Some(os);
            }
            session_data.info.is_elevated = is_elevated;
            session_data.info.uptime_secs = uptime_secs;
            session_data.info.display_count = display_count;
        }
    }

    /// Update heartbeat timestamp
    pub async fn update_heartbeat(&self, session_id: SessionId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.get_mut(&session_id) {
            session_data.info.last_heartbeat = chrono::Utc::now();
            session_data.last_heartbeat_instant = Instant::now();
        }
    }

    /// Check if a session has timed out (no heartbeat for too long)
    pub async fn is_session_timed_out(&self, session_id: SessionId) -> bool {
        let sessions = self.sessions.read().await;
        if let Some(session_data) = sessions.get(&session_id) {
            session_data.last_heartbeat_instant.elapsed().as_secs() > HEARTBEAT_TIMEOUT_SECS
        } else {
            true // Non-existent sessions are considered timed out
        }
    }

    /// Get sessions that have timed out
    pub async fn get_timed_out_sessions(&self) -> Vec<SessionId> {
        let sessions = self.sessions.read().await;
        sessions
            .iter()
            .filter(|(_, data)| data.last_heartbeat_instant.elapsed().as_secs() > HEARTBEAT_TIMEOUT_SECS)
            .map(|(id, _)| *id)
            .collect()
    }

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

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

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

    /// Join a session as a viewer, returns channels and sends StartStream to agent
    pub async fn join_session(&self, session_id: SessionId, viewer_id: ViewerId, viewer_name: String) -> Option<(FrameReceiver, InputSender)> {
        let mut sessions = self.sessions.write().await;
        let session_data = sessions.get_mut(&session_id)?;

        let was_empty = session_data.viewers.is_empty();

        // Add viewer info
        let viewer_info = ViewerInfo {
            id: viewer_id.clone(),
            name: viewer_name.clone(),
            connected_at: chrono::Utc::now(),
        };
        session_data.viewers.insert(viewer_id.clone(), viewer_info);

        // Update session info
        session_data.info.viewer_count = session_data.viewers.len();
        session_data.info.viewers = session_data.viewers.values().cloned().collect();

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

        // If this is the first viewer, send StartStream to agent
        if was_empty {
            tracing::info!("Viewer {} ({}) joined session {}, sending StartStream", viewer_name, viewer_id, session_id);
            Self::send_start_stream_internal(session_data, &viewer_id).await;
        } else {
            tracing::info!("Viewer {} ({}) joined session {}", viewer_name, viewer_id, session_id);
        }

        Some((frame_rx, input_tx))
    }

    /// Internal helper to send StartStream message
    async fn send_start_stream_internal(session_data: &SessionData, viewer_id: &str) {
        use crate::proto;
        use prost::Message;

        let start_stream = proto::Message {
            payload: Some(proto::message::Payload::StartStream(proto::StartStream {
                viewer_id: viewer_id.to_string(),
                display_id: 0, // Primary display
            })),
        };

        let mut buf = Vec::new();
        if start_stream.encode(&mut buf).is_ok() {
            let _ = session_data.input_tx.send(buf).await;
        }
    }

    /// Leave a session as a viewer, sends StopStream if no viewers left
    pub async fn leave_session(&self, session_id: SessionId, viewer_id: &ViewerId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.get_mut(&session_id) {
            let viewer_name = session_data.viewers.get(viewer_id).map(|v| v.name.clone());
            session_data.viewers.remove(viewer_id);
            session_data.info.viewer_count = session_data.viewers.len();
            session_data.info.viewers = session_data.viewers.values().cloned().collect();

            // If no more viewers, send StopStream to agent
            if session_data.viewers.is_empty() {
                tracing::info!("Last viewer {} ({}) left session {}, sending StopStream",
                    viewer_name.as_deref().unwrap_or("unknown"), viewer_id, session_id);
                Self::send_stop_stream_internal(session_data, viewer_id).await;
            } else {
                tracing::info!("Viewer {} ({}) left session {}",
                    viewer_name.as_deref().unwrap_or("unknown"), viewer_id, session_id);
            }
        }
    }

    /// Internal helper to send StopStream message
    async fn send_stop_stream_internal(session_data: &SessionData, viewer_id: &str) {
        use crate::proto;
        use prost::Message;

        let stop_stream = proto::Message {
            payload: Some(proto::message::Payload::StopStream(proto::StopStream {
                viewer_id: viewer_id.to_string(),
            })),
        };

        let mut buf = Vec::new();
        if stop_stream.encode(&mut buf).is_ok() {
            let _ = session_data.input_tx.send(buf).await;
        }
    }

    /// Mark agent as disconnected
    /// For persistent agents: keep session but mark as offline
    /// For support sessions: remove session entirely
    pub async fn mark_agent_disconnected(&self, session_id: SessionId) {
        let mut sessions = self.sessions.write().await;
        if let Some(session_data) = sessions.get_mut(&session_id) {
            if session_data.info.is_persistent {
                // Persistent agent - keep session but mark as offline
                tracing::info!("Persistent agent {} marked offline (session {} preserved)",
                    session_data.info.agent_id, session_id);
                session_data.info.is_online = false;
                session_data.info.is_streaming = false;
                session_data.info.viewer_count = 0;
                session_data.info.viewers.clear();
                session_data.viewers.clear();
            } else {
                // Support session - remove entirely
                let agent_id = session_data.info.agent_id.clone();
                sessions.remove(&session_id);
                drop(sessions); // Release sessions lock before acquiring agents lock
                let mut agents = self.agents.write().await;
                agents.remove(&agent_id);
                tracing::info!("Support session {} removed", session_id);
            }
        }
    }

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

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

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

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

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

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

+    /// Restore a machine as an offline session (called on startup from database)

+    pub async fn restore_offline_machine(&self, agent_id: &str, hostname: &str) -> SessionId {

+        let session_id = Uuid::new_v4();

+            agent_id: agent_id.to_string(),

+            agent_name: hostname.to_string(),

+            is_online: false,  // Offline until agent reconnects

+        // Create placeholder channels (will be replaced on reconnect)

+        let (frame_tx, _) = broadcast::channel(16);

+        let (input_tx, input_rx) = tokio::sync::mpsc::channel(64);

+        let session_data = SessionData {

+            last_heartbeat_instant: Instant::now(),

+        let mut sessions = self.sessions.write().await;

+        sessions.insert(session_id, session_data);

+        let mut agents = self.agents.write().await;

+        agents.insert(agent_id.to_string(), session_id);

+        tracing::info!("Restored offline machine: {} ({})", hostname, agent_id);


Good. Now I need to update the relay/mod.rs to add database logging for events. Let me read the current relay/mod.rs file and add the database operations.This is the most complex update - I need to:1. Pass the database to the handlers2. Log events on connect/disconnect/viewer join/leave
3. Upsert machines on connect4. Mark machines offline on disconnectThis file is already loaded in the context from earlier. Let me make the changes.


Now let me update the relay handlers to add database logging. First, let me check the current state of the file:


     1→//! WebSocket relay handlers
     2→//!
     3→//! Handles WebSocket connections from agents and viewers,
     4→//! relaying video frames and input events between them.
     5→
     6→use axum::{
     7→    extract::{
     8→        ws::{Message, WebSocket, WebSocketUpgrade},
     9→        Query, State,
    10→    },
    11→    response::IntoResponse,
    12→};
    13→use futures_util::{SinkExt, StreamExt};
    14→use prost::Message as ProstMessage;
    15→use serde::Deserialize;
    16→use tracing::{error, info, warn};
    17→use uuid::Uuid;
    18→
    19→use crate::proto;
    20→use crate::session::SessionManager;
    21→use crate::AppState;
    22→
    23→#[derive(Debug, Deserialize)]
    24→pub struct AgentParams {
    25→    agent_id: String,
    26→    #[serde(default)]
    27→    agent_name: Option<String>,
    28→    #[serde(default)]
    29→    support_code: Option<String>,
    30→    #[serde(default)]
    31→    hostname: Option<String>,
    32→}
    33→
    34→#[derive(Debug, Deserialize)]
    35→pub struct ViewerParams {
    36→    session_id: String,
    37→    #[serde(default = "default_viewer_name")]
    38→    viewer_name: String,
    39→}
    40→
    41→fn default_viewer_name() -> String {
    42→    "Technician".to_string()
    43→}
    44→
    45→/// WebSocket handler for agent connections
    46→pub async fn agent_ws_handler(
    47→    ws: WebSocketUpgrade,
    48→    State(state): State<AppState>,
    49→    Query(params): Query<AgentParams>,
    50→) -> impl IntoResponse {
    51→    let agent_id = params.agent_id;
    52→    let agent_name = params.hostname.or(params.agent_name).unwrap_or_else(|| agent_id.clone());
    53→    let support_code = params.support_code;
    54→    let sessions = state.sessions.clone();
    55→    let support_codes = state.support_codes.clone();
    56→
    57→    ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, agent_id, agent_name, support_code))
    58→}
    59→
    60→/// WebSocket handler for viewer connections
    61→pub async fn viewer_ws_handler(
    62→    ws: WebSocketUpgrade,
    63→    State(state): State<AppState>,
    64→    Query(params): Query<ViewerParams>,
    65→) -> impl IntoResponse {
    66→    let session_id = params.session_id;
    67→    let viewer_name = params.viewer_name;
    68→    let sessions = state.sessions.clone();
    69→
    70→    ws.on_upgrade(move |socket| handle_viewer_connection(socket, sessions, session_id, viewer_name))
    71→}
    72→
    73→/// Handle an agent WebSocket connection
    74→async fn handle_agent_connection(
    75→    socket: WebSocket,
    76→    sessions: SessionManager,
    77→    support_codes: crate::support_codes::SupportCodeManager,
    78→    agent_id: String,
    79→    agent_name: String,
    80→    support_code: Option<String>,

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

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

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

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

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

#[derive(Debug, Deserialize)]
pub struct ViewerParams {
    session_id: String,
    #[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();

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

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

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

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


I need to update the relay handlers to pass the database and log events. Let me update the agent and viewer WebSocket handlers to include the db parameter.


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

use crate::proto;
use crate::session::SessionManager;
use crate::db::{self, Database};
use crate::AppState;


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:
    15→use serde::Deserialize;
    16→use tracing::{error, info, warn};
    17→use uuid::Uuid;
    18→
    19→use crate::proto;
    20→use crate::session::SessionManager;
    21→use crate::db::{self, Database};
    22→use crate::AppState;
    23→
    24→#[derive(Debug, Deserialize)]
    25→pub struct AgentParams {
    26→    agent_id: String,

use crate::proto;
use crate::session::SessionManager;
use crate::db::{self, Database};
use crate::AppState;

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

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

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

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

#[derive(Debug, Deserialize)]
pub struct ViewerParams {
    session_id: String,
    #[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();

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

use crate::proto;
use crate::session::SessionManager;
+use crate::db::{self, Database};
use crate::AppState;
#[derive(Debug, Deserialize)]


Good. Now let me update the agent_ws_handler to pass the db:


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

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

/// WebSocket handler for 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))
}


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:
    42→fn default_viewer_name() -> String {
    43→    "Technician".to_string()
    44→}
    45→
    46→/// WebSocket handler for agent connections
    47→pub async fn agent_ws_handler(
    48→    ws: WebSocketUpgrade,
    49→    State(state): State<AppState>,
    50→    Query(params): Query<AgentParams>,
    51→) -> impl IntoResponse {
    52→    let agent_id = params.agent_id;
    53→    let agent_name = params.hostname.or(params.agent_name).unwrap_or_else(|| agent_id.clone());
    54→    let support_code = params.support_code;
    55→    let sessions = state.sessions.clone();
    56→    let support_codes = state.support_codes.clone();
    57→    let db = state.db.clone();
    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,

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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

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

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