From 22f592dd27ec688008eb669dc3751a9236e80892 Mon Sep 17 00:00:00 2001 From: Mike Swanson Date: Thu, 1 Jan 2026 10:40:11 -0700 Subject: [PATCH] Add organization/site/tags support for machine grouping - Added organization, site, tags columns to connect_machines table - Agent now sends org/site/tags from embedded config in AgentStatus - Server stores org/site/tags metadata in database - Enables grouping machines by client/site/tag in dashboard --- agent/src/session/mod.rs | 3 +++ proto/guruconnect.proto | 3 +++ server/src/db/machines.rs | 31 +++++++++++++++++++++++++++++++ server/src/relay/mod.rs | 28 ++++++++++++++++++++++++++-- server/src/session/mod.rs | 21 +++++++++++++++++++++ 5 files changed, 84 insertions(+), 2 deletions(-) diff --git a/agent/src/session/mod.rs b/agent/src/session/mod.rs index e92ab18..4573035 100644 --- a/agent/src/session/mod.rs +++ b/agent/src/session/mod.rs @@ -175,6 +175,9 @@ impl SessionManager { display_count: self.get_display_count(), is_streaming: self.state == SessionState::Streaming, agent_version: crate::build_info::short_version(), + organization: self.config.company.clone().unwrap_or_default(), + site: self.config.site.clone().unwrap_or_default(), + tags: self.config.tags.clone(), }; let msg = Message { diff --git a/proto/guruconnect.proto b/proto/guruconnect.proto index 1079b4f..cf974e0 100644 --- a/proto/guruconnect.proto +++ b/proto/guruconnect.proto @@ -277,6 +277,9 @@ message AgentStatus { int32 display_count = 5; bool is_streaming = 6; string agent_version = 7; // Agent version (e.g., "0.1.0-abc123") + string organization = 8; // Company/organization name + string site = 9; // Site/location name + repeated string tags = 10; // Tags for categorization } // Server commands agent to uninstall itself diff --git a/server/src/db/machines.rs b/server/src/db/machines.rs index 41caaef..fd44a9b 100644 --- a/server/src/db/machines.rs +++ b/server/src/db/machines.rs @@ -116,3 +116,34 @@ pub async fn delete_machine(pool: &PgPool, agent_id: &str) -> Result<(), sqlx::E .await?; Ok(()) } + +/// Update machine organization, site, and tags +pub async fn update_machine_metadata( + pool: &PgPool, + agent_id: &str, + organization: Option<&str>, + site: Option<&str>, + tags: &[String], +) -> Result<(), sqlx::Error> { + // Only update if at least one value is provided + if organization.is_none() && site.is_none() && tags.is_empty() { + return Ok(()); + } + + sqlx::query( + r#" + UPDATE connect_machines SET + organization = COALESCE($1, organization), + site = COALESCE($2, site), + tags = CASE WHEN $3::text[] = '{}' THEN tags ELSE $3 END + WHERE agent_id = $4 + "#, + ) + .bind(organization) + .bind(site) + .bind(tags) + .bind(agent_id) + .execute(pool) + .await?; + Ok(()) +} diff --git a/server/src/relay/mod.rs b/server/src/relay/mod.rs index 58099a6..b7719a5 100644 --- a/server/src/relay/mod.rs +++ b/server/src/relay/mod.rs @@ -313,6 +313,16 @@ async fn handle_agent_connection( } else { Some(status.agent_version.clone()) }; + let organization = if status.organization.is_empty() { + None + } else { + Some(status.organization.clone()) + }; + let site = if status.site.is_empty() { + None + } else { + Some(status.site.clone()) + }; sessions_status.update_agent_status( session_id, Some(status.os_version.clone()), @@ -321,6 +331,9 @@ async fn handle_agent_connection( status.display_count, status.is_streaming, agent_version.clone(), + organization.clone(), + site.clone(), + status.tags.clone(), ).await; // Update version in database if present @@ -328,8 +341,19 @@ async fn handle_agent_connection( let _ = crate::db::releases::update_machine_version(db.pool(), &agent_id, version).await; } - info!("Agent status update: {} - streaming={}, uptime={}s, version={:?}", - status.hostname, status.is_streaming, status.uptime_secs, agent_version); + // Update organization/site/tags in database if present + if let Some(ref db) = db { + let _ = crate::db::machines::update_machine_metadata( + db.pool(), + &agent_id, + organization.as_deref(), + site.as_deref(), + &status.tags, + ).await; + } + + info!("Agent status update: {} - streaming={}, uptime={}s, version={:?}, org={:?}, site={:?}", + status.hostname, status.is_streaming, status.uptime_secs, agent_version, organization, site); } Some(proto::message::Payload::Heartbeat(_)) => { // Update heartbeat timestamp diff --git a/server/src/session/mod.rs b/server/src/session/mod.rs index 9d88c67..64cc804 100644 --- a/server/src/session/mod.rs +++ b/server/src/session/mod.rs @@ -48,6 +48,9 @@ pub struct Session { pub uptime_secs: i64, pub display_count: i32, pub agent_version: Option, // Agent software version + pub organization: Option, // Company/organization name + pub site: Option, // Site/location name + pub tags: Vec, // Tags for categorization } /// Channel for sending frames from agent to viewers @@ -140,6 +143,9 @@ impl SessionManager { uptime_secs: 0, display_count: 1, agent_version: None, + organization: None, + site: None, + tags: Vec::new(), }; let session_data = SessionData { @@ -170,6 +176,9 @@ impl SessionManager { display_count: i32, is_streaming: bool, agent_version: Option, + organization: Option, + site: Option, + tags: Vec, ) { let mut sessions = self.sessions.write().await; if let Some(session_data) = sessions.get_mut(&session_id) { @@ -185,6 +194,15 @@ impl SessionManager { if let Some(version) = agent_version { session_data.info.agent_version = Some(version); } + if let Some(org) = organization { + session_data.info.organization = Some(org); + } + if let Some(s) = site { + session_data.info.site = Some(s); + } + if !tags.is_empty() { + session_data.info.tags = tags; + } } } @@ -461,6 +479,9 @@ impl SessionManager { uptime_secs: 0, display_count: 1, agent_version: None, + organization: None, + site: None, + tags: Vec::new(), }; // Create placeholder channels (will be replaced on reconnect)