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
This commit is contained in:
2026-01-01 10:40:11 -07:00
parent 5a82637a04
commit 22f592dd27
5 changed files with 84 additions and 2 deletions

View File

@@ -175,6 +175,9 @@ impl SessionManager {
display_count: self.get_display_count(), display_count: self.get_display_count(),
is_streaming: self.state == SessionState::Streaming, is_streaming: self.state == SessionState::Streaming,
agent_version: crate::build_info::short_version(), 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 { let msg = Message {

View File

@@ -277,6 +277,9 @@ message AgentStatus {
int32 display_count = 5; int32 display_count = 5;
bool is_streaming = 6; bool is_streaming = 6;
string agent_version = 7; // Agent version (e.g., "0.1.0-abc123") 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 // Server commands agent to uninstall itself

View File

@@ -116,3 +116,34 @@ pub async fn delete_machine(pool: &PgPool, agent_id: &str) -> Result<(), sqlx::E
.await?; .await?;
Ok(()) 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(())
}

View File

@@ -313,6 +313,16 @@ async fn handle_agent_connection(
} else { } else {
Some(status.agent_version.clone()) 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( sessions_status.update_agent_status(
session_id, session_id,
Some(status.os_version.clone()), Some(status.os_version.clone()),
@@ -321,6 +331,9 @@ async fn handle_agent_connection(
status.display_count, status.display_count,
status.is_streaming, status.is_streaming,
agent_version.clone(), agent_version.clone(),
organization.clone(),
site.clone(),
status.tags.clone(),
).await; ).await;
// Update version in database if present // 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; let _ = crate::db::releases::update_machine_version(db.pool(), &agent_id, version).await;
} }
info!("Agent status update: {} - streaming={}, uptime={}s, version={:?}", // Update organization/site/tags in database if present
status.hostname, status.is_streaming, status.uptime_secs, agent_version); 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(_)) => { Some(proto::message::Payload::Heartbeat(_)) => {
// Update heartbeat timestamp // Update heartbeat timestamp

View File

@@ -48,6 +48,9 @@ pub struct Session {
pub uptime_secs: i64, pub uptime_secs: i64,
pub display_count: i32, pub display_count: i32,
pub agent_version: Option<String>, // Agent software version pub agent_version: Option<String>, // Agent software version
pub organization: Option<String>, // Company/organization name
pub site: Option<String>, // Site/location name
pub tags: Vec<String>, // Tags for categorization
} }
/// Channel for sending frames from agent to viewers /// Channel for sending frames from agent to viewers
@@ -140,6 +143,9 @@ impl SessionManager {
uptime_secs: 0, uptime_secs: 0,
display_count: 1, display_count: 1,
agent_version: None, agent_version: None,
organization: None,
site: None,
tags: Vec::new(),
}; };
let session_data = SessionData { let session_data = SessionData {
@@ -170,6 +176,9 @@ impl SessionManager {
display_count: i32, display_count: i32,
is_streaming: bool, is_streaming: bool,
agent_version: Option<String>, agent_version: Option<String>,
organization: Option<String>,
site: Option<String>,
tags: Vec<String>,
) { ) {
let mut sessions = self.sessions.write().await; let mut sessions = self.sessions.write().await;
if let Some(session_data) = sessions.get_mut(&session_id) { if let Some(session_data) = sessions.get_mut(&session_id) {
@@ -185,6 +194,15 @@ impl SessionManager {
if let Some(version) = agent_version { if let Some(version) = agent_version {
session_data.info.agent_version = Some(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, uptime_secs: 0,
display_count: 1, display_count: 1,
agent_version: None, agent_version: None,
organization: None,
site: None,
tags: Vec::new(),
}; };
// Create placeholder channels (will be replaced on reconnect) // Create placeholder channels (will be replaced on reconnect)