From 654aa3be7456dfdf7ffdc11a4ece26066d95b707 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 02:00:13 +0800 Subject: [PATCH 01/44] refactor(server): own the service monitor check transition in one module Both the background scheduler and the manual HTTP trigger previously assembled the check transition themselves, with three behavioral gaps: record and state were written in two separate statements, slow checks for the same monitor could overlap, and the manual path skipped the maintenance gate and notifications entirely (a manual success after failures cleared last_status, permanently losing the recovery notification). Move the complete transition behind MonitorCheckRunner in service::monitor_check: per-monitor in-flight guard, checker dispatch, transactional record+state commit, maintenance gating, and failure/recovery notifications. The scheduler and the HTTP handler now only decide when a check happens; a concurrent manual trigger returns 409 Conflict instead of interleaving writes. --- .../server/src/router/api/service_monitor.rs | 46 +- crates/server/src/service/mod.rs | 1 + crates/server/src/service/monitor_check.rs | 408 ++++++++++++++++++ crates/server/src/service/service_monitor.rs | 14 +- crates/server/src/state.rs | 6 + .../src/task/service_monitor_checker.rs | 274 +----------- 6 files changed, 456 insertions(+), 293 deletions(-) create mode 100644 crates/server/src/service/monitor_check.rs diff --git a/crates/server/src/router/api/service_monitor.rs b/crates/server/src/router/api/service_monitor.rs index 4d198eeb2..ca493bb31 100644 --- a/crates/server/src/router/api/service_monitor.rs +++ b/crates/server/src/router/api/service_monitor.rs @@ -8,7 +8,7 @@ use serde::{Deserialize, Serialize}; use crate::entity::{service_monitor, service_monitor_record}; use crate::error::{ApiResponse, AppError, ok}; -use crate::service::checker; +use crate::service::monitor_check::CheckOutcome; use crate::service::service_monitor::{ CreateServiceMonitor, ServiceMonitorService, UpdateServiceMonitor, }; @@ -204,6 +204,7 @@ async fn get_records( responses( (status = 200, description = "Check result", body = service_monitor_record::Model), (status = 404, description = "Not found"), + (status = 409, description = "A check for this monitor is already running"), ), security(("session_cookie" = []), ("api_key" = []), ("bearer_token" = [])) )] @@ -211,37 +212,14 @@ async fn trigger_check( State(state): State>, Path(id): Path, ) -> Result>, AppError> { - let monitor = ServiceMonitorService::get(&state.db, &id).await?; - - let config: serde_json::Value = serde_json::from_str(&monitor.config_json).unwrap_or_default(); - - let result = checker::run_check(&monitor.monitor_type, &monitor.target, &config).await; - - // Insert the record - let record = ServiceMonitorService::insert_record( - &state.db, - &monitor.id, - result.success, - result.latency, - result.detail, - result.error, - ) - .await?; - - // Update monitor state - let consecutive_failures = if result.success { - 0 - } else { - monitor.consecutive_failures + 1 - }; - - ServiceMonitorService::update_check_state( - &state.db, - &monitor.id, - record.success, - consecutive_failures, - ) - .await?; - - ok(record) + match state + .monitor_check_runner + .run_check(&state.db, &state.config, &id) + .await? + { + CheckOutcome::Completed(record) => ok(record), + CheckOutcome::AlreadyRunning => Err(AppError::Conflict(format!( + "A check for monitor {id} is already running" + ))), + } } diff --git a/crates/server/src/service/mod.rs b/crates/server/src/service/mod.rs index a51165261..0c10b63a8 100644 --- a/crates/server/src/service/mod.rs +++ b/crates/server/src/service/mod.rs @@ -22,6 +22,7 @@ pub mod ip_quality; pub mod ip_risk; pub mod maintenance; pub mod mobile_auth; +pub mod monitor_check; pub mod network_probe; pub mod notification; pub mod oauth; diff --git a/crates/server/src/service/monitor_check.rs b/crates/server/src/service/monitor_check.rs new file mode 100644 index 000000000..94c4379d4 --- /dev/null +++ b/crates/server/src/service/monitor_check.rs @@ -0,0 +1,408 @@ +//! The complete service-monitor check transition. +//! +//! Both the background scheduler ([`crate::task::service_monitor_checker`]) +//! and the manual HTTP trigger execute checks through +//! [`MonitorCheckRunner::run_check`], which owns the entire transition: +//! per-monitor overlap guard, checker dispatch, transactional record + state +//! write, maintenance gating, and failure/recovery notifications. Callers +//! only choose *when* a check happens; *what a check means* lives here. + +use chrono::Utc; +use dashmap::DashMap; +use sea_orm::{DatabaseConnection, TransactionTrait}; + +use crate::config::AppConfig; +use crate::entity::{service_monitor, service_monitor_record}; +use crate::error::AppError; +use crate::service::checker::{self, CheckResult}; +use crate::service::maintenance::MaintenanceService; +use crate::service::notification::{NotificationService, NotifyContext}; +use crate::service::service_monitor::ServiceMonitorService; + +/// Outcome of a check request. +pub enum CheckOutcome { + /// The check ran to completion; the inserted record is returned. + Completed(service_monitor_record::Model), + /// A check for this monitor is still in flight; nothing was run. + AlreadyRunning, +} + +/// Owns the check transition for all callers and guarantees at most one +/// in-flight check per monitor. +pub struct MonitorCheckRunner { + in_flight: DashMap, +} + +impl Default for MonitorCheckRunner { + fn default() -> Self { + Self::new() + } +} + +impl MonitorCheckRunner { + pub fn new() -> Self { + Self { + in_flight: DashMap::new(), + } + } + + /// Run one complete check transition for the monitor with `monitor_id`. + /// + /// Returns [`CheckOutcome::AlreadyRunning`] without touching the database + /// when a check for the same monitor has not finished yet, so slow checks + /// can never interleave their record/state writes. Notification failures + /// are logged, not returned: the check itself succeeded. + pub async fn run_check( + &self, + db: &DatabaseConnection, + config: &AppConfig, + monitor_id: &str, + ) -> Result { + let Some(_guard) = self.try_acquire(monitor_id) else { + return Ok(CheckOutcome::AlreadyRunning); + }; + + // Load fresh state so a run dispatched from a stale model (the + // scheduler fetches monitors up to a tick earlier) still sees the + // current failure count and last status. + let monitor = ServiceMonitorService::get(db, monitor_id).await?; + + let checker_config: serde_json::Value = + serde_json::from_str(&monitor.config_json).unwrap_or_default(); + let result = checker::run_check(&monitor.monitor_type, &monitor.target, &checker_config).await; + + let consecutive_failures = if result.success { + 0 + } else { + monitor.consecutive_failures + 1 + }; + + // The record and the monitor state describe the same transition — + // commit them together so readers never see one without the other. + let txn = db.begin().await?; + let record = ServiceMonitorService::insert_record( + &txn, + &monitor.id, + result.success, + result.latency, + result.detail.clone(), + result.error.clone(), + ) + .await?; + ServiceMonitorService::update_check_state( + &txn, + &monitor.id, + result.success, + consecutive_failures, + ) + .await?; + txn.commit().await?; + + self.notify(db, config, &monitor, &result, consecutive_failures) + .await; + + Ok(CheckOutcome::Completed(record)) + } + + /// Reserve the in-flight slot for `monitor_id`, or return `None` if a + /// check is already running. The slot is released when the guard drops, + /// including on panic or early return. + fn try_acquire(&self, monitor_id: &str) -> Option> { + match self.in_flight.entry(monitor_id.to_string()) { + dashmap::mapref::entry::Entry::Occupied(_) => None, + dashmap::mapref::entry::Entry::Vacant(v) => { + v.insert(()); + Some(InFlightGuard { + map: &self.in_flight, + id: monitor_id.to_string(), + }) + } + } + } + + /// Maintenance gate plus failure/recovery notifications. + /// + /// `monitor` carries the *pre-check* state (`last_status`, + /// `consecutive_failures`), which decides whether this transition is a + /// new failure past the retry threshold or a recovery. + async fn notify( + &self, + db: &DatabaseConnection, + config: &AppConfig, + monitor: &service_monitor::Model, + result: &CheckResult, + consecutive_failures: i32, + ) { + let Some(ref group_id) = monitor.notification_group_id else { + return; + }; + + let failure_crossed_threshold = + !result.success && consecutive_failures > monitor.retry_count; + let recovered = result.success && monitor.last_status == Some(false); + if !failure_crossed_threshold && !recovered { + return; + } + + if self.any_server_in_maintenance(db, monitor).await { + tracing::debug!( + "Skipping notification for service monitor '{}': associated server in maintenance", + monitor.name + ); + return; + } + + let ctx = if failure_crossed_threshold { + let error_msg = result.error.as_deref().unwrap_or("Unknown error"); + NotifyContext { + server_name: monitor.name.clone(), + server_id: monitor.id.clone(), + rule_name: format!("{} ({})", monitor.name, monitor.monitor_type), + event: "triggered".to_string(), + message: format!( + "Service monitor '{}' failed after {} consecutive failures: {}", + monitor.name, consecutive_failures, error_msg + ), + time: Utc::now().format("%Y-%m-%d %H:%M:%S UTC").to_string(), + ..Default::default() + } + } else { + NotifyContext { + server_name: monitor.name.clone(), + server_id: monitor.id.clone(), + rule_name: format!("{} ({})", monitor.name, monitor.monitor_type), + event: "recovered".to_string(), + message: format!( + "Service monitor '{}' has recovered after {} consecutive failures", + monitor.name, monitor.consecutive_failures + ), + time: Utc::now().format("%Y-%m-%d %H:%M:%S UTC").to_string(), + ..Default::default() + } + }; + + if let Err(e) = NotificationService::send_group(db, config, group_id, &ctx).await { + tracing::error!( + "Failed to send {} notification for {}: {e}", + ctx.event, + monitor.name + ); + } + } + + async fn any_server_in_maintenance( + &self, + db: &DatabaseConnection, + monitor: &service_monitor::Model, + ) -> bool { + let Some(ref server_ids_json) = monitor.server_ids_json else { + return false; + }; + let server_ids: Vec = serde_json::from_str(server_ids_json).unwrap_or_default(); + for sid in &server_ids { + if MaintenanceService::is_in_maintenance(db, sid) + .await + .unwrap_or(false) + { + return true; + } + } + false + } +} + +/// Releases the per-monitor in-flight slot on drop. +struct InFlightGuard<'a> { + map: &'a DashMap, + id: String, +} + +impl Drop for InFlightGuard<'_> { + fn drop(&mut self) { + self.map.remove(&self.id); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::entity::service_monitor; + use crate::test_utils::setup_test_db; + use chrono::TimeZone; + use sea_orm::{ActiveModelTrait, Set}; + + fn fixed_now() -> chrono::DateTime { + Utc.with_ymd_and_hms(2026, 6, 25, 12, 0, 0).unwrap() + } + + /// Insert a monitor row so `run_check` can write a record and update state. + async fn insert_monitor( + db: &DatabaseConnection, + id: &str, + monitor_type: &str, + target: &str, + retry_count: i32, + last_status: Option, + consecutive_failures: i32, + ) -> service_monitor::Model { + service_monitor::ActiveModel { + id: Set(id.to_string()), + name: Set(format!("monitor-{id}")), + monitor_type: Set(monitor_type.to_string()), + target: Set(target.to_string()), + interval: Set(60), + config_json: Set("{}".to_string()), + notification_group_id: Set(None), + retry_count: Set(retry_count), + server_ids_json: Set(None), + enabled: Set(true), + last_status: Set(last_status), + consecutive_failures: Set(consecutive_failures), + last_checked_at: Set(None), + created_at: Set(fixed_now()), + updated_at: Set(fixed_now()), + } + .insert(db) + .await + .expect("insert monitor should succeed") + } + + #[tokio::test] + async fn run_check_writes_failure_record_for_offline_target() { + let (db, _tmp) = setup_test_db().await; + // Invalid TCP target ("no-port") fails fast and offline (no network). + insert_monitor(&db, "mon-fail", "tcp", "no-port-here", 1, None, 0).await; + let runner = MonitorCheckRunner::new(); + + let outcome = runner + .run_check(&db, &AppConfig::default(), "mon-fail") + .await + .unwrap(); + + // The outcome carries the failure record that was written. + let CheckOutcome::Completed(record) = outcome else { + panic!("expected Completed outcome"); + }; + assert!(!record.success); + assert!( + record.error.is_some(), + "offline check should record an error message" + ); + + let records = ServiceMonitorService::get_records(&db, "mon-fail", None, None, None) + .await + .unwrap(); + assert_eq!(records.len(), 1); + + // Monitor state advanced: last_status=false, consecutive_failures bumped to 1. + let updated = ServiceMonitorService::get(&db, "mon-fail").await.unwrap(); + assert_eq!(updated.last_status, Some(false)); + assert_eq!(updated.consecutive_failures, 1); + assert!(updated.last_checked_at.is_some()); + } + + #[tokio::test] + async fn run_check_accumulates_consecutive_failures() { + let (db, _tmp) = setup_test_db().await; + // Seed a monitor that has already failed once. + insert_monitor(&db, "mon-fail2", "tcp", "no-port-here", 1, Some(false), 1).await; + let runner = MonitorCheckRunner::new(); + + runner + .run_check(&db, &AppConfig::default(), "mon-fail2") + .await + .unwrap(); + + // The failure count increments from the monitor's prior value (1 -> 2). + let updated = ServiceMonitorService::get(&db, "mon-fail2").await.unwrap(); + assert_eq!(updated.consecutive_failures, 2); + assert_eq!(updated.last_status, Some(false)); + } + + #[tokio::test] + async fn run_check_handles_unknown_monitor_type() { + let (db, _tmp) = setup_test_db().await; + // Unknown type returns a deterministic failure without any network call. + insert_monitor(&db, "mon-unknown", "bogus", "whatever", 1, None, 0).await; + let runner = MonitorCheckRunner::new(); + + runner + .run_check(&db, &AppConfig::default(), "mon-unknown") + .await + .unwrap(); + + let records = ServiceMonitorService::get_records(&db, "mon-unknown", None, None, None) + .await + .unwrap(); + assert_eq!(records.len(), 1); + assert!(!records[0].success); + assert!( + records[0] + .error + .as_deref() + .unwrap_or_default() + .contains("Unknown monitor type") + ); + } + + #[tokio::test] + async fn run_check_unknown_monitor_id_is_not_found() { + let (db, _tmp) = setup_test_db().await; + let runner = MonitorCheckRunner::new(); + + let result = runner + .run_check(&db, &AppConfig::default(), "missing") + .await; + assert!(matches!(result, Err(AppError::NotFound(_)))); + } + + #[tokio::test] + async fn run_check_skips_when_already_in_flight() { + let (db, _tmp) = setup_test_db().await; + insert_monitor(&db, "mon-busy", "tcp", "no-port-here", 1, None, 0).await; + let runner = MonitorCheckRunner::new(); + + // Simulate a check that has not finished yet. + let guard = runner.try_acquire("mon-busy").expect("slot must be free"); + + let outcome = runner + .run_check(&db, &AppConfig::default(), "mon-busy") + .await + .unwrap(); + assert!(matches!(outcome, CheckOutcome::AlreadyRunning)); + + // No record was written and state did not advance. + let records = ServiceMonitorService::get_records(&db, "mon-busy", None, None, None) + .await + .unwrap(); + assert!(records.is_empty()); + + // Once the in-flight check finishes, the next run proceeds. + drop(guard); + let outcome = runner + .run_check(&db, &AppConfig::default(), "mon-busy") + .await + .unwrap(); + assert!(matches!(outcome, CheckOutcome::Completed(_))); + } + + #[tokio::test] + async fn in_flight_slot_is_released_after_completion() { + let (db, _tmp) = setup_test_db().await; + insert_monitor(&db, "mon-seq", "tcp", "no-port-here", 1, None, 0).await; + let runner = MonitorCheckRunner::new(); + + for _ in 0..2 { + let outcome = runner + .run_check(&db, &AppConfig::default(), "mon-seq") + .await + .unwrap(); + assert!(matches!(outcome, CheckOutcome::Completed(_))); + } + + let records = ServiceMonitorService::get_records(&db, "mon-seq", None, None, None) + .await + .unwrap(); + assert_eq!(records.len(), 2); + } +} diff --git a/crates/server/src/service/service_monitor.rs b/crates/server/src/service/service_monitor.rs index 43e456b82..1fc650dff 100644 --- a/crates/server/src/service/service_monitor.rs +++ b/crates/server/src/service/service_monitor.rs @@ -298,8 +298,11 @@ impl ServiceMonitorService { } /// Update runtime state columns after a check completes. - pub async fn update_check_state( - db: &DatabaseConnection, + /// + /// Generic over the connection so `monitor_check` can commit it in the + /// same transaction as the record insert. + pub async fn update_check_state( + db: &C, id: &str, success: bool, consecutive_failures: i32, @@ -322,8 +325,11 @@ impl ServiceMonitorService { } /// Insert a check result record. - pub async fn insert_record( - db: &DatabaseConnection, + /// + /// Generic over the connection so `monitor_check` can commit it in the + /// same transaction as the state update. + pub async fn insert_record( + db: &C, monitor_id: &str, success: bool, latency: Option, diff --git a/crates/server/src/state.rs b/crates/server/src/state.rs index 68a084800..6e837a6df 100644 --- a/crates/server/src/state.rs +++ b/crates/server/src/state.rs @@ -17,6 +17,7 @@ use crate::service::geoip::GeoIpService; use crate::service::high_risk_audit::{ DockerLogsAuditContext, ExecAuditContext, TerminalAuditContext, }; +use crate::service::monitor_check::MonitorCheckRunner; use crate::service::firewall::FirewallService; use crate::service::security::SecurityService; use crate::service::task_scheduler::TaskScheduler; @@ -114,6 +115,10 @@ pub struct AppState { pub exec_audit_contexts: DashMap, /// DNS PTR enricher for traceroute hops (shared across requests). pub traceroute_enricher: crate::service::traceroute_enrich::TracerouteEnricher, + /// Owns the service-monitor check transition (overlap guard, transactional + /// record/state write, maintenance gate, notifications) for both the + /// scheduler and the manual HTTP trigger. + pub monitor_check_runner: MonitorCheckRunner, } static RATE_CHECK_COUNTER: AtomicU64 = AtomicU64::new(0); @@ -264,6 +269,7 @@ impl AppState { docker_logs_audit_contexts: DashMap::new(), exec_audit_contexts: DashMap::new(), traceroute_enricher, + monitor_check_runner: MonitorCheckRunner::new(), })) } } diff --git a/crates/server/src/task/service_monitor_checker.rs b/crates/server/src/task/service_monitor_checker.rs index f6ced9be6..eb7f3a9a0 100644 --- a/crates/server/src/task/service_monitor_checker.rs +++ b/crates/server/src/task/service_monitor_checker.rs @@ -5,9 +5,7 @@ use std::time::Instant; use chrono::Utc; use tokio::sync::Semaphore; -use crate::service::checker; -use crate::service::maintenance::MaintenanceService; -use crate::service::notification::{NotificationService, NotifyContext}; +use crate::service::monitor_check::CheckOutcome; use crate::service::service_monitor::ServiceMonitorService; use crate::state::AppState; @@ -62,7 +60,20 @@ pub async fn run(state: Arc) { Err(_) => return, }; - execute_check(&state, &monitor).await; + match state + .monitor_check_runner + .run_check(&state.db, &state.config, &monitor.id) + .await + { + Ok(CheckOutcome::Completed(_)) => {} + Ok(CheckOutcome::AlreadyRunning) => tracing::debug!( + "Skipping check for '{}': previous check still running", + monitor.name + ), + Err(e) => { + tracing::error!("Service monitor check failed for {}: {e}", monitor.id); + } + } }); } } @@ -131,151 +142,13 @@ pub(crate) fn select_due_monitors( due_monitors } -/// Execute a single monitor check: run the checker, insert the record, update state, -/// and send notifications if needed. -pub(crate) async fn execute_check( - state: &AppState, - monitor: &crate::entity::service_monitor::Model, -) { - let config: serde_json::Value = serde_json::from_str(&monitor.config_json).unwrap_or_default(); - - let result = checker::run_check(&monitor.monitor_type, &monitor.target, &config).await; - - // Insert the record - if let Err(e) = ServiceMonitorService::insert_record( - &state.db, - &monitor.id, - result.success, - result.latency, - result.detail.clone(), - result.error.clone(), - ) - .await - { - tracing::error!( - "Failed to insert service monitor record for {}: {e}", - monitor.id - ); - return; - } - - // Calculate new consecutive failure count - let consecutive_failures = if result.success { - 0 - } else { - monitor.consecutive_failures + 1 - }; - - // Update monitor state - if let Err(e) = ServiceMonitorService::update_check_state( - &state.db, - &monitor.id, - result.success, - consecutive_failures, - ) - .await - { - tracing::error!("Failed to update check state for {}: {e}", monitor.id); - return; - } - - // Determine if we need to send notifications - let was_failing = monitor.last_status == Some(false); - - // Skip notifications if any associated server is in maintenance - let in_maintenance = if let Some(ref server_ids_json) = monitor.server_ids_json { - let server_ids: Vec = serde_json::from_str(server_ids_json).unwrap_or_default(); - let mut any_in_maintenance = false; - for sid in &server_ids { - if MaintenanceService::is_in_maintenance(&state.db, sid) - .await - .unwrap_or(false) - { - any_in_maintenance = true; - break; - } - } - any_in_maintenance - } else { - false - }; - - if in_maintenance { - tracing::debug!( - "Skipping notification for service monitor '{}': associated server in maintenance", - monitor.name - ); - return; - } - - // Failure notification: consecutive failures exceeded retry_count - if !result.success - && consecutive_failures > monitor.retry_count - && let Some(ref group_id) = monitor.notification_group_id - { - let error_msg = result.error.as_deref().unwrap_or("Unknown error"); - let ctx = NotifyContext { - server_name: monitor.name.clone(), - server_id: monitor.id.clone(), - rule_name: format!("{} ({})", monitor.name, monitor.monitor_type), - event: "triggered".to_string(), - message: format!( - "Service monitor '{}' failed after {} consecutive failures: {}", - monitor.name, consecutive_failures, error_msg - ), - time: Utc::now().format("%Y-%m-%d %H:%M:%S UTC").to_string(), - ..Default::default() - }; - - if let Err(e) = - NotificationService::send_group(&state.db, &state.config, group_id, &ctx).await - { - tracing::error!( - "Failed to send failure notification for {}: {e}", - monitor.name - ); - } - } - - // Recovery notification: was failing, now succeeded - if result.success - && was_failing - && let Some(ref group_id) = monitor.notification_group_id - { - let ctx = NotifyContext { - server_name: monitor.name.clone(), - server_id: monitor.id.clone(), - rule_name: format!("{} ({})", monitor.name, monitor.monitor_type), - event: "recovered".to_string(), - message: format!( - "Service monitor '{}' has recovered after {} consecutive failures", - monitor.name, monitor.consecutive_failures - ), - time: Utc::now().format("%Y-%m-%d %H:%M:%S UTC").to_string(), - ..Default::default() - }; - - if let Err(e) = - NotificationService::send_group(&state.db, &state.config, group_id, &ctx).await - { - tracing::error!( - "Failed to send recovery notification for {}: {e}", - monitor.name - ); - } - } -} - #[cfg(test)] mod tests { - // `super::*` already brings `ServiceMonitorService`, `HashMap`, `Instant`, - // `Utc`, and `AppState` into scope from the parent module. + // `super::*` already brings `HashMap`, `Instant`, and `Utc` into scope + // from the parent module. use super::*; - use crate::config::AppConfig; use crate::entity::service_monitor; - use crate::test_utils::setup_test_db; use chrono::TimeZone; - use sea_orm::{ActiveModelTrait, DatabaseConnection, Set}; /// Fixed wall-clock instant used so scheduling math is deterministic. fn fixed_now() -> chrono::DateTime { @@ -307,38 +180,6 @@ mod tests { } } - /// Insert a monitor row so `execute_check` can write a record and update state. - async fn insert_monitor( - db: &DatabaseConnection, - id: &str, - monitor_type: &str, - target: &str, - retry_count: i32, - last_status: Option, - consecutive_failures: i32, - ) -> service_monitor::Model { - service_monitor::ActiveModel { - id: Set(id.to_string()), - name: Set(format!("monitor-{id}")), - monitor_type: Set(monitor_type.to_string()), - target: Set(target.to_string()), - interval: Set(60), - config_json: Set("{}".to_string()), - notification_group_id: Set(None), - retry_count: Set(retry_count), - server_ids_json: Set(None), - enabled: Set(true), - last_status: Set(last_status), - consecutive_failures: Set(consecutive_failures), - last_checked_at: Set(None), - created_at: Set(fixed_now()), - updated_at: Set(fixed_now()), - } - .insert(db) - .await - .expect("insert monitor should succeed") - } - // ---- select_due_monitors (pure scheduler core) ---- #[test] @@ -428,83 +269,6 @@ mod tests { assert!(due.is_empty()); } - // ---- execute_check (full DB side-effect path) ---- - - #[tokio::test] - async fn execute_check_writes_failure_record_for_offline_target() { - let (db, _tmp) = setup_test_db().await; - // Invalid TCP target ("no-port") fails fast and offline (no network). - insert_monitor(&db, "mon-fail", "tcp", "no-port-here", 1, None, 0).await; - let state = AppState::new(db, AppConfig::default()).await.unwrap(); - - let monitor = ServiceMonitorService::get(&state.db, "mon-fail") - .await - .unwrap(); - execute_check(&state, &monitor).await; - - // A failure record was written. - let records = ServiceMonitorService::get_records(&state.db, "mon-fail", None, None, None) - .await - .unwrap(); - assert_eq!(records.len(), 1); - assert!(!records[0].success); - assert!( - records[0].error.is_some(), - "offline check should record an error message" - ); - - // Monitor state advanced: last_status=false, consecutive_failures bumped to 1. - let updated = ServiceMonitorService::get(&state.db, "mon-fail") - .await - .unwrap(); - assert_eq!(updated.last_status, Some(false)); - assert_eq!(updated.consecutive_failures, 1); - assert!(updated.last_checked_at.is_some()); - } - - #[tokio::test] - async fn execute_check_accumulates_consecutive_failures() { - let (db, _tmp) = setup_test_db().await; - // Seed a monitor that has already failed once. - insert_monitor(&db, "mon-fail2", "tcp", "no-port-here", 1, Some(false), 1).await; - let state = AppState::new(db, AppConfig::default()).await.unwrap(); - - let monitor = ServiceMonitorService::get(&state.db, "mon-fail2") - .await - .unwrap(); - execute_check(&state, &monitor).await; - - // The failure count increments from the monitor's prior value (1 -> 2). - let updated = ServiceMonitorService::get(&state.db, "mon-fail2") - .await - .unwrap(); - assert_eq!(updated.consecutive_failures, 2); - assert_eq!(updated.last_status, Some(false)); - } - - #[tokio::test] - async fn execute_check_handles_unknown_monitor_type() { - let (db, _tmp) = setup_test_db().await; - // Unknown type returns a deterministic failure without any network call. - insert_monitor(&db, "mon-unknown", "bogus", "whatever", 1, None, 0).await; - let state = AppState::new(db, AppConfig::default()).await.unwrap(); - - let monitor = ServiceMonitorService::get(&state.db, "mon-unknown") - .await - .unwrap(); - execute_check(&state, &monitor).await; - - let records = ServiceMonitorService::get_records(&state.db, "mon-unknown", None, None, None) - .await - .unwrap(); - assert_eq!(records.len(), 1); - assert!(!records[0].success); - assert!( - records[0] - .error - .as_deref() - .unwrap_or_default() - .contains("Unknown monitor type") - ); - } + // The full check transition (record/state writes, overlap guard, + // notifications) is owned and tested by `service::monitor_check`. } From 8b0dbbadbb8ab6896de839c1e520ef55598addea Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 02:12:06 +0800 Subject: [PATCH 02/44] refactor(server): share one credential policy across HTTP and WebSocket auth The browser, terminal, and docker-logs WebSocket handlers each carried a private copy of the credential ladder (cookie -> API key -> bearer) plus header extractors, duplicating precedence, forced-password, session-source, and mobile-expiry rules that middleware::auth already claimed to own. Any future change to one rule could silently drift between the four copies. Deepen middleware::auth with resolve_connection, returning identity plus lifetime facts (AuthenticatedConnection { user, mobile_expires }), and a resolve_ws_connection view that rejects must_change_password users because onboarding cannot complete over a WebSocket. All three WS adapters and the HTTP middleware now resolve through the same seam while keeping their protocol-specific responses; api/auth.rs change_password reuses the shared raw-token extractors. Two edge cases now follow one policy instead of per-copy accidents: a flagged user's valid cookie no longer falls through to a weaker credential on WS upgrades, and a non-web session enforces its fixed expiry whichever header carried it. --- crates/server/src/middleware/auth.rs | 228 ++++++++++++++++++--- crates/server/src/router/api/auth.rs | 29 +-- crates/server/src/router/ws/browser.rs | 90 +------- crates/server/src/router/ws/docker_logs.rs | 90 +------- crates/server/src/router/ws/terminal.rs | 90 +------- 5 files changed, 232 insertions(+), 295 deletions(-) diff --git a/crates/server/src/middleware/auth.rs b/crates/server/src/middleware/auth.rs index 8ebf9d79a..3df9931cd 100644 --- a/crates/server/src/middleware/auth.rs +++ b/crates/server/src/middleware/auth.rs @@ -55,28 +55,48 @@ fn is_onboarding_whitelisted(method: &axum::http::Method, path: &str) -> bool { ) } -/// Resolve the authenticated user (if any) from a request's headers. +/// Identity plus lifetime facts for an authenticated connection. /// -/// Tries, in order: session cookie, `X-API-Key` header, `Bearer` token. -/// Returns `None` when no credential is present or none validates. +/// Long-lived transports (WebSockets) need more than the user: they must know +/// when a fixed-lifetime credential expires so they can close the connection +/// mid-stream. Short-lived HTTP requests simply ignore `mobile_expires`. +#[derive(Debug, Clone)] +pub struct AuthenticatedConnection { + pub user: CurrentUser, + /// Fixed expiry of the authenticating session when it is a non-web + /// (mobile) session, whatever header carried it. Web sessions renew on + /// use (sliding expiry) and API keys never expire, so both yield `None`. + pub mobile_expires: Option>, +} + +/// Resolve the authenticated connection (if any) from a request's headers. /// -/// This is the single, shared credential-parsing routine. Both -/// `auth_middleware` (which enforces auth) and public routes that need -/// optional auth (e.g. the public status page's IP masking) call it, so the -/// security-sensitive parsing logic never drifts between copies. -pub async fn resolve_optional_user( +/// Tries, in order: session cookie, `X-API-Key` header, `Bearer` token. The +/// first credential that validates decides the identity and lifetime; later +/// rungs are only consulted when earlier ones are absent or invalid. Returns +/// `None` when no credential validates. +/// +/// This is the single, shared credential policy. HTTP middleware, public +/// routes with optional auth, and every WebSocket handler resolve through it, +/// so precedence, session-source, and expiry semantics never drift between +/// copies. +pub async fn resolve_connection( headers: &HeaderMap, state: &AppState, -) -> Option { +) -> Option { + let session_ttl = state.config.auth.session_ttl; + // Try session cookie if let Some(token) = extract_session_cookie(headers) - && let Some((user, _session)) = - AuthService::validate_session(&state.db, &token, state.config.auth.session_ttl) - .await - .ok() - .flatten() + && let Some((user, session)) = AuthService::validate_session(&state.db, &token, session_ttl) + .await + .ok() + .flatten() { - return Some(CurrentUser::from(user)); + return Some(AuthenticatedConnection { + user: CurrentUser::from(user), + mobile_expires: mobile_expiry(&session), + }); } // Try API key header @@ -86,23 +106,56 @@ pub async fn resolve_optional_user( .ok() .flatten() { - return Some(CurrentUser::from(user)); + return Some(AuthenticatedConnection { + user: CurrentUser::from(user), + mobile_expires: None, + }); } // Try Bearer token if let Some(token) = extract_bearer_token(headers) - && let Some((user, _session)) = - AuthService::validate_session(&state.db, &token, state.config.auth.session_ttl) - .await - .ok() - .flatten() + && let Some((user, session)) = AuthService::validate_session(&state.db, &token, session_ttl) + .await + .ok() + .flatten() { - return Some(CurrentUser::from(user)); + return Some(AuthenticatedConnection { + user: CurrentUser::from(user), + mobile_expires: mobile_expiry(&session), + }); } None } +/// A non-web session has a fixed lifetime (no sliding renewal), so its expiry +/// is a fact the transport must enforce; web sessions yield `None`. +fn mobile_expiry(session: &crate::entity::session::Model) -> Option> { + (session.source != "web").then_some(session.expires_at) +} + +/// Resolve the authenticated user for a WebSocket upgrade. +/// +/// Same credential policy as [`resolve_connection`], plus the WS-specific +/// rule: a user flagged `must_change_password` is rejected outright, because +/// the onboarding flow cannot be completed over a WebSocket. +pub async fn resolve_ws_connection( + headers: &HeaderMap, + state: &AppState, +) -> Option { + resolve_connection(headers, state) + .await + .filter(|conn| !conn.user.must_change_password) +} + +/// Resolve the authenticated user (if any) from a request's headers. +/// +/// Identity-only view of [`resolve_connection`] for callers that do not care +/// about connection lifetime (HTTP middleware, optional-auth public routes). +pub async fn resolve_optional_user(headers: &HeaderMap, state: &AppState) -> Option { + resolve_connection(headers, state).await.map(|c| c.user) +} + pub async fn auth_middleware( State(state): State>, mut req: Request, @@ -140,7 +193,8 @@ pub async fn require_admin(req: Request, next: Next) -> Response { next.run(req).await } -fn extract_session_cookie(headers: &HeaderMap) -> Option { +/// Extract the `session_token` value from the Cookie header, if present. +pub(crate) fn extract_session_cookie(headers: &HeaderMap) -> Option { headers .get("cookie")? .to_str() @@ -152,7 +206,8 @@ fn extract_session_cookie(headers: &HeaderMap) -> Option { }) } -fn extract_api_key(headers: &HeaderMap) -> Option { +/// Extract the raw API key from the `X-API-Key` header, if present. +pub(crate) fn extract_api_key(headers: &HeaderMap) -> Option { headers .get("x-api-key")? .to_str() @@ -160,7 +215,8 @@ fn extract_api_key(headers: &HeaderMap) -> Option { .map(|s| s.to_string()) } -fn extract_bearer_token(headers: &HeaderMap) -> Option { +/// Extract the bearer token from the Authorization header, if present. +pub(crate) fn extract_bearer_token(headers: &HeaderMap) -> Option { headers .get("authorization")? .to_str() @@ -256,4 +312,126 @@ mod tests { assert!(!is_onboarding_whitelisted(&Method::GET, "/servers")); assert!(!is_onboarding_whitelisted(&Method::GET, "/api/auth/me")); } + + // ── resolve_connection / resolve_ws_connection (shared credential policy) ── + + mod resolve { + use super::*; + use crate::config::AppConfig; + use crate::entity::{session, user}; + use crate::service::auth::AuthService; + use crate::test_utils::setup_test_db; + use chrono::Utc; + use sea_orm::{ActiveModelTrait, DatabaseConnection, Set}; + + /// Seed a session row and return its plaintext token. + async fn seed_session( + db: &DatabaseConnection, + user_id: &str, + source: &str, + expires_at: chrono::DateTime, + ) -> String { + let token = AuthService::generate_session_token(); + session::ActiveModel { + id: Set(uuid::Uuid::new_v4().to_string()), + user_id: Set(user_id.to_string()), + token: Set(AuthService::hash_session_token(&token)), + ip: Set("127.0.0.1".into()), + user_agent: Set("test".into()), + expires_at: Set(expires_at), + created_at: Set(Utc::now()), + source: Set(source.to_string()), + mobile_session_id: Set(None), + } + .insert(db) + .await + .expect("seed session"); + token + } + + #[tokio::test] + async fn cookie_takes_precedence_over_bearer() { + let (db, _tmp) = setup_test_db().await; + let cookie_user = AuthService::create_user(&db, "cookie-user", "pass1234", "admin") + .await + .unwrap(); + let bearer_user = AuthService::create_user(&db, "bearer-user", "pass1234", "member") + .await + .unwrap(); + let far = Utc::now() + chrono::Duration::hours(1); + let cookie_tok = seed_session(&db, &cookie_user.id, "web", far).await; + let bearer_tok = seed_session(&db, &bearer_user.id, "web", far).await; + let state = AppState::new(db, AppConfig::default()).await.unwrap(); + + let mut headers = HeaderMap::new(); + headers.insert( + "cookie", + HeaderValue::from_str(&format!("session_token={cookie_tok}")).unwrap(), + ); + headers.insert( + "authorization", + HeaderValue::from_str(&format!("Bearer {bearer_tok}")).unwrap(), + ); + + let conn = resolve_connection(&headers, &state).await.unwrap(); + assert_eq!(conn.user.username, "cookie-user"); + } + + #[tokio::test] + async fn mobile_bearer_session_carries_fixed_expiry() { + let (db, _tmp) = setup_test_db().await; + let user = AuthService::create_user(&db, "mob", "pass1234", "member") + .await + .unwrap(); + let fixed = Utc::now() + chrono::Duration::hours(24); + let token = seed_session(&db, &user.id, "mobile", fixed).await; + let state = AppState::new(db, AppConfig::default()).await.unwrap(); + + let headers = headers_with("authorization", &format!("Bearer {token}")); + let conn = resolve_connection(&headers, &state).await.unwrap(); + assert_eq!( + conn.mobile_expires.map(|e| e.timestamp()), + Some(fixed.timestamp()), + "non-web session must expose its fixed expiry" + ); + } + + #[tokio::test] + async fn web_session_has_no_fixed_expiry() { + let (db, _tmp) = setup_test_db().await; + let user = AuthService::create_user(&db, "webby", "pass1234", "member") + .await + .unwrap(); + let token = + seed_session(&db, &user.id, "web", Utc::now() + chrono::Duration::hours(1)).await; + let state = AppState::new(db, AppConfig::default()).await.unwrap(); + + let headers = headers_with( + "cookie", + &format!("session_token={token}"), + ); + let conn = resolve_connection(&headers, &state).await.unwrap(); + assert!(conn.mobile_expires.is_none()); + } + + #[tokio::test] + async fn ws_rejects_must_change_password_user() { + let (db, _tmp) = setup_test_db().await; + let user = AuthService::create_user(&db, "flagged", "pass1234", "admin") + .await + .unwrap(); + let mut active: user::ActiveModel = user.clone().into(); + active.must_change_password = Set(true); + active.update(&db).await.unwrap(); + let token = + seed_session(&db, &user.id, "web", Utc::now() + chrono::Duration::hours(1)).await; + let state = AppState::new(db, AppConfig::default()).await.unwrap(); + + let headers = headers_with("cookie", &format!("session_token={token}")); + // HTTP resolution still authenticates (middleware whitelists the + // onboarding routes); the WS view rejects outright. + assert!(resolve_connection(&headers, &state).await.is_some()); + assert!(resolve_ws_connection(&headers, &state).await.is_none()); + } + } } diff --git a/crates/server/src/router/api/auth.rs b/crates/server/src/router/api/auth.rs index b28fc2a38..0cc0d57be 100644 --- a/crates/server/src/router/api/auth.rs +++ b/crates/server/src/router/api/auth.rs @@ -437,8 +437,8 @@ pub async fn change_password( // previously stolen session. An API-key authenticated caller has no session // token here, so `current_token` is None and all of the user's sessions are // revoked — the API key itself is unaffected. - let current_token = extract_session_cookie_token(&req_headers) - .or_else(|| extract_bearer_token_value(&req_headers)); + let current_token = crate::middleware::auth::extract_session_cookie(&req_headers) + .or_else(|| crate::middleware::auth::extract_bearer_token(&req_headers)); AuthService::change_password( &state.db, ¤t_user.user_id, @@ -712,31 +712,6 @@ fn extract_user_agent(headers: &HeaderMap) -> String { .to_string() } -/// Extract the `session_token` value from the Cookie header, if present. -fn extract_session_cookie_token(headers: &HeaderMap) -> Option { - headers - .get("cookie")? - .to_str() - .ok()? - .split(';') - .find_map(|cookie| { - cookie - .trim() - .strip_prefix("session_token=") - .map(|v| v.to_string()) - }) -} - -/// Extract the bearer token from the Authorization header, if present. -fn extract_bearer_token_value(headers: &HeaderMap) -> Option { - headers - .get("authorization")? - .to_str() - .ok()? - .strip_prefix("Bearer ") - .map(|s| s.to_string()) -} - /// Build the `detail` string for login audit rows (failed / rate-limited). /// The username is attacker-controlled on the public login endpoints, so it is /// truncated to a sane bound before being persisted, keeping audit rows from diff --git a/crates/server/src/router/ws/browser.rs b/crates/server/src/router/ws/browser.rs index 02dcd3e3d..f2c905c62 100644 --- a/crates/server/src/router/ws/browser.rs +++ b/crates/server/src/router/ws/browser.rs @@ -11,8 +11,8 @@ use futures_util::{SinkExt, StreamExt}; use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, QueryOrder}; use crate::entity::{agent_enrollment, server_tag}; +use crate::middleware::auth::resolve_ws_connection; use crate::service::agent_manager::aggregate_disk_io; -use crate::service::auth::AuthService; use crate::service::server::ServerService; use crate::state::AppState; use serverbee_common::constants::MAX_WS_MESSAGE_SIZE; @@ -28,88 +28,20 @@ async fn browser_ws_handler( headers: HeaderMap, ws: WebSocketUpgrade, ) -> Response { - // Validate auth: try session cookie first, then API key, then Bearer token - let auth = validate_browser_auth(&state, &headers).await; - match auth { - Some((_user_id, is_admin, mobile_expires)) => ws - .max_message_size(MAX_WS_MESSAGE_SIZE) - .on_upgrade(move |socket| handle_browser_ws(socket, state, is_admin, mobile_expires)), + // Shared credential policy (precedence, session source, mobile expiry) + // lives in `middleware::auth`; this adapter only maps the verdict onto + // the browser WS protocol. + match resolve_ws_connection(&headers, &state).await { + Some(conn) => { + let is_admin = conn.user.role == "admin"; + ws.max_message_size(MAX_WS_MESSAGE_SIZE).on_upgrade(move |socket| { + handle_browser_ws(socket, state, is_admin, conn.mobile_expires) + }) + } None => axum::http::StatusCode::UNAUTHORIZED.into_response(), } } -/// Returns `Some((user_id, mobile_expires))` on success. -/// `mobile_expires` is `Some(expires_at)` when authenticated via a non-web session -/// (Bearer token from mobile), so the WS connection can be auto-closed on expiry. -/// For web sessions and API keys, `mobile_expires` is `None` (they use sliding expiry -/// or never expire respectively). -async fn validate_browser_auth( - state: &Arc, - headers: &HeaderMap, -) -> Option<(String, bool, Option>)> { - // Try session cookie (always web source → no mobile expiry) - if let Some(token) = extract_session_cookie(headers) - && let Ok(Some((user, _session))) = - AuthService::validate_session(&state.db, &token, state.config.auth.session_ttl).await - && !user.must_change_password - { - return Some((user.id, user.role == "admin", None)); - } - - // Try API key header (no expiry) - if let Some(key) = extract_api_key(headers) - && let Ok(Some(user)) = AuthService::validate_api_key(&state.db, &key).await - && !user.must_change_password - { - return Some((user.id, user.role == "admin", None)); - } - - // Try Bearer token (may be a mobile session with a fixed expiry) - if let Some(token) = extract_bearer_token(headers) - && let Ok(Some((user, session))) = - AuthService::validate_session(&state.db, &token, state.config.auth.session_ttl).await - && !user.must_change_password - { - let mobile_expires = if session.source != "web" { - Some(session.expires_at) - } else { - None - }; - return Some((user.id, user.role == "admin", mobile_expires)); - } - - None -} - -fn extract_session_cookie(headers: &HeaderMap) -> Option { - headers - .get("cookie")? - .to_str() - .ok()? - .split(';') - .find_map(|cookie| { - let cookie = cookie.trim(); - cookie.strip_prefix("session_token=").map(|v| v.to_string()) - }) -} - -fn extract_api_key(headers: &HeaderMap) -> Option { - headers - .get("x-api-key")? - .to_str() - .ok() - .map(|s| s.to_string()) -} - -fn extract_bearer_token(headers: &HeaderMap) -> Option { - headers - .get("authorization")? - .to_str() - .ok()? - .strip_prefix("Bearer ") - .map(|s| s.to_string()) -} - async fn handle_browser_ws( socket: WebSocket, state: Arc, diff --git a/crates/server/src/router/ws/docker_logs.rs b/crates/server/src/router/ws/docker_logs.rs index 94640f6e0..db78c67e8 100644 --- a/crates/server/src/router/ws/docker_logs.rs +++ b/crates/server/src/router/ws/docker_logs.rs @@ -10,6 +10,7 @@ use futures_util::{SinkExt, StreamExt}; use serde::Deserialize; use tokio::sync::mpsc; +use crate::middleware::auth::resolve_ws_connection; use crate::router::utils::extract_client_ip; use crate::service::audit::AuditService; use crate::service::high_risk_audit::DockerLogsAuditContext; @@ -35,14 +36,15 @@ async fn docker_logs_ws_handler( ) .to_string(); - // Auth: session cookie or API key - let user = validate_auth(&state, &headers).await; - match user { - Some((user_id, role, mobile_expires)) => { + // Shared credential policy lives in `middleware::auth`; this adapter only + // applies Docker-specific rules (admin role, agent online, capability). + match resolve_ws_connection(&headers, &state).await { + Some(conn) => { + let user_id = conn.user.user_id; // Docker log streaming exposes sensitive container output // (env vars, connection strings, tokens), so it is admin-only, // consistent with terminal access. - if role != "admin" { + if conn.user.role != "admin" { let detail = serde_json::json!({ "server_id": server_id, "deny_reason": "role_forbidden", @@ -77,89 +79,13 @@ async fn docker_logs_ws_handler( } ws.max_message_size(MAX_WS_MESSAGE_SIZE) .on_upgrade(move |socket| { - handle_docker_logs_ws(socket, state, server_id, user_id, ip, mobile_expires) + handle_docker_logs_ws(socket, state, server_id, user_id, ip, conn.mobile_expires) }) } None => axum::http::StatusCode::UNAUTHORIZED.into_response(), } } -/// Returns `(user_id, role, mobile_expires)` on success. -/// -/// `mobile_expires` is `Some(expires_at)` only for a non-web (mobile) Bearer -/// session, whose lifetime is fixed (no sliding renewal). The handler uses it to -/// force-close the WS when the token expires mid-session, matching the browser -/// WS. Web sessions (sliding expiry) and API keys never expire the socket and -/// yield `None`. -async fn validate_auth( - state: &Arc, - headers: &HeaderMap, -) -> Option<(String, String, Option>)> { - use crate::service::auth::AuthService; - - // Try session cookie (always web source → no mobile expiry) - if let Some(token) = extract_session_cookie(headers) - && let Ok(Some((user, _session))) = - AuthService::validate_session(&state.db, &token, state.config.auth.session_ttl).await - && !user.must_change_password - { - return Some((user.id, user.role, None)); - } - - // Try API key header (no expiry) - if let Some(key) = extract_api_key(headers) - && let Ok(Some(user)) = AuthService::validate_api_key(&state.db, &key).await - && !user.must_change_password - { - return Some((user.id, user.role, None)); - } - - // Try Bearer token (may be a mobile session with a fixed expiry) - if let Some(token) = extract_bearer_token(headers) - && let Ok(Some((user, session))) = - AuthService::validate_session(&state.db, &token, state.config.auth.session_ttl).await - && !user.must_change_password - { - let mobile_expires = if session.source != "web" { - Some(session.expires_at) - } else { - None - }; - return Some((user.id, user.role, mobile_expires)); - } - - None -} - -fn extract_session_cookie(headers: &HeaderMap) -> Option { - headers - .get("cookie")? - .to_str() - .ok()? - .split(';') - .find_map(|cookie| { - let cookie = cookie.trim(); - cookie.strip_prefix("session_token=").map(|v| v.to_string()) - }) -} - -fn extract_api_key(headers: &HeaderMap) -> Option { - headers - .get("x-api-key")? - .to_str() - .ok() - .map(|s| s.to_string()) -} - -fn extract_bearer_token(headers: &HeaderMap) -> Option { - headers - .get("authorization")? - .to_str() - .ok()? - .strip_prefix("Bearer ") - .map(|s| s.to_string()) -} - /// Browser -> Server messages for docker logs #[derive(Deserialize)] #[serde(tag = "type", rename_all = "snake_case")] diff --git a/crates/server/src/router/ws/terminal.rs b/crates/server/src/router/ws/terminal.rs index f0b003454..fdb0da14b 100644 --- a/crates/server/src/router/ws/terminal.rs +++ b/crates/server/src/router/ws/terminal.rs @@ -10,6 +10,7 @@ use futures_util::{SinkExt, StreamExt}; use serde::Deserialize; use tokio::sync::mpsc; +use crate::middleware::auth::resolve_ws_connection; use crate::router::utils::extract_client_ip; use crate::service::agent_manager::TerminalSessionEvent; use crate::service::audit::AuditService; @@ -36,12 +37,13 @@ async fn terminal_ws_handler( ) .to_string(); - // Auth: session cookie or API key - let user = validate_auth(&state, &headers).await; - match user { - Some((user_id, role, mobile_expires)) => { + // Shared credential policy lives in `middleware::auth`; this adapter only + // applies terminal-specific rules (admin role, agent online, capability). + match resolve_ws_connection(&headers, &state).await { + Some(conn) => { + let user_id = conn.user.user_id; // Terminal access is admin-only - if role != "admin" { + if conn.user.role != "admin" { let detail = serde_json::json!({ "server_id": server_id, "deny_reason": "role_forbidden", @@ -75,89 +77,13 @@ async fn terminal_ws_handler( return error.into_response(); } ws.max_message_size(MAX_WS_MESSAGE_SIZE).on_upgrade(move |socket| { - handle_terminal_ws(socket, state, server_id, user_id, ip, mobile_expires) + handle_terminal_ws(socket, state, server_id, user_id, ip, conn.mobile_expires) }) } None => axum::http::StatusCode::UNAUTHORIZED.into_response(), } } -/// Returns `(user_id, role, mobile_expires)` on success. -/// -/// `mobile_expires` is `Some(expires_at)` only for a non-web (mobile) Bearer -/// session, whose lifetime is fixed (no sliding renewal). The handler uses it to -/// force-close the WS when the token expires mid-session, matching the browser -/// WS. Web sessions (sliding expiry) and API keys never expire the socket and -/// yield `None`. -async fn validate_auth( - state: &Arc, - headers: &HeaderMap, -) -> Option<(String, String, Option>)> { - use crate::service::auth::AuthService; - - // Try session cookie (always web source → no mobile expiry) - if let Some(token) = extract_session_cookie(headers) - && let Ok(Some((user, _session))) = - AuthService::validate_session(&state.db, &token, state.config.auth.session_ttl).await - && !user.must_change_password - { - return Some((user.id, user.role, None)); - } - - // Try API key header (no expiry) - if let Some(key) = extract_api_key(headers) - && let Ok(Some(user)) = AuthService::validate_api_key(&state.db, &key).await - && !user.must_change_password - { - return Some((user.id, user.role, None)); - } - - // Try Bearer token (may be a mobile session with a fixed expiry) - if let Some(token) = extract_bearer_token(headers) - && let Ok(Some((user, session))) = - AuthService::validate_session(&state.db, &token, state.config.auth.session_ttl).await - && !user.must_change_password - { - let mobile_expires = if session.source != "web" { - Some(session.expires_at) - } else { - None - }; - return Some((user.id, user.role, mobile_expires)); - } - - None -} - -fn extract_session_cookie(headers: &HeaderMap) -> Option { - headers - .get("cookie")? - .to_str() - .ok()? - .split(';') - .find_map(|cookie| { - let cookie = cookie.trim(); - cookie.strip_prefix("session_token=").map(|v| v.to_string()) - }) -} - -fn extract_api_key(headers: &HeaderMap) -> Option { - headers - .get("x-api-key")? - .to_str() - .ok() - .map(|s| s.to_string()) -} - -fn extract_bearer_token(headers: &HeaderMap) -> Option { - headers - .get("authorization")? - .to_str() - .ok()? - .strip_prefix("Bearer ") - .map(|s| s.to_string()) -} - /// Browser terminal WS message format (JSON) #[derive(Deserialize)] #[serde(tag = "type", rename_all = "snake_case")] From 8a6d372505c0cb8db68acdbf2de631977ee0e483 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 03:08:52 +0800 Subject: [PATCH 03/44] refactor(agent): own capability state and lifecycle in one runtime authority MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Effective capabilities were a per-connection Arc handed to every subsystem, with a per-connection grant supervisor rewriting it. Consumers knew the cell, the bit timing, and who updates when — and two lifecycle gaps followed from that shape: a temporary grant's expiry never reconciled running work (live terminal sessions survived revocation), and the SecurityManager was started once at boot from the base bitmask, so a temporary security-events grant could never start it (nor could a persisted grant enable it across a restart) and expiry never stopped it. Introduce a process-wide CapabilityAuthority (capability_grants::authority) that owns the base bitmask, folds in grants, and drives every transition. It exposes has()/effective()/active_grants(), a state watch for long-running subsystems, and a transition broadcast that connections forward to the server as CapabilitiesChanged. All consumers (terminal, pinger, network prober, docker, file manager, IP-quality checker, reporter dispatch) now gate through the authority instead of sharing the atomic. Lifecycle fixes riding on the seam: - the reporter's transition arm reconciles in-flight work before announcing: losing CAP_TERMINAL tears down live PTY sessions with a TerminalError - SecurityManager::spawn_supervised starts/stops the pipeline as CAP_SECURITY_EVENTS becomes effective or expires, at any point in the agent's lifetime The old per-connection grant supervisor is deleted; its pure evaluate() logic and tests move into the authority. An e2e test covers the full expiry path: grant folded into connect-time SystemInfo, live session torn down on expiry, CapabilitiesChanged announced without the bit. --- .../agent/src/capability_grants/authority.rs | 362 ++++++++++++++++++ crates/agent/src/capability_grants/mod.rs | 4 +- .../agent/src/capability_grants/supervisor.rs | 307 --------------- crates/agent/src/docker/mod.rs | 15 +- crates/agent/src/file_manager.rs | 13 +- crates/agent/src/ip_quality/mod.rs | 23 +- crates/agent/src/main.rs | 30 +- crates/agent/src/network_prober.rs | 34 +- crates/agent/src/pinger.rs | 37 +- crates/agent/src/reporter/file_ops.rs | 11 +- crates/agent/src/reporter/mod.rs | 274 +++++++++---- crates/agent/src/security/manager.rs | 40 ++ crates/agent/src/terminal.rs | 53 +-- 13 files changed, 714 insertions(+), 489 deletions(-) create mode 100644 crates/agent/src/capability_grants/authority.rs delete mode 100644 crates/agent/src/capability_grants/supervisor.rs diff --git a/crates/agent/src/capability_grants/authority.rs b/crates/agent/src/capability_grants/authority.rs new file mode 100644 index 000000000..3440868ba --- /dev/null +++ b/crates/agent/src/capability_grants/authority.rs @@ -0,0 +1,362 @@ +//! Agent-owned authority over effective runtime capabilities. +//! +//! One process-wide [`CapabilityAuthority`] owns the base bitmask, folds in +//! temporary grants from the grants file, and drives every transition: +//! updating the effective state, notifying long-running subsystems so they can +//! reconcile in-flight work, and handing connections the change events they +//! forward to the server. Consumers ask [`CapabilityAuthority::has`] / +//! [`CapabilityAuthority::effective`] and never learn where the bits come +//! from or when they change. + +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::{AtomicU32, Ordering}; +use std::time::Duration; + +use tokio::sync::{broadcast, watch}; + +use serverbee_common::constants::{ALL_CAPABILITIES, CAP_VALID_MASK, has_capability}; +use serverbee_common::protocol::{CapabilityChangeAction, CapabilityChangeEvent, TemporaryGrant}; + +use super::store::CapabilityGrantStore; + +/// One transition of the effective capability state, as observed by the +/// authority's evaluation loop. +#[derive(Debug, Clone)] +pub struct CapabilityTransition { + /// The new effective bitmask. + pub effective: u32, + /// Temporary grants active after the transition. + pub temporary: Vec, + /// Per-capability change events (granted / expired / revoked). + pub changes: Vec, +} + +pub struct CapabilityAuthority { + base: u32, + grants_path: PathBuf, + effective: AtomicU32, + state_tx: watch::Sender, + transition_tx: broadcast::Sender, +} + +impl CapabilityAuthority { + /// Build the authority, seeding the effective state from `base` plus any + /// still-active grants in the grants file (so grants survive a restart). + /// Call [`Self::run`] once to start the transition loop. + pub fn new(base: u32, grants_path: PathBuf) -> Arc { + let store = CapabilityGrantStore::load(&grants_path); + let effective = (base | store.active_bits(now_unix(), base)) & CAP_VALID_MASK; + let (state_tx, _) = watch::channel(effective); + let (transition_tx, _) = broadcast::channel(16); + Arc::new(Self { + base, + grants_path, + effective: AtomicU32::new(effective), + state_tx, + transition_tx, + }) + } + + /// Fixed-state authority whose effective caps equal `base` and never + /// change (no grants file, no running loop). Test-only: production code + /// always gates on the process-wide authority built in `main`. + #[cfg(test)] + pub fn fixed(base: u32) -> Arc { + Self::new(base, PathBuf::from("/nonexistent/capability_grants.json")) + } + + /// Whether the capability bit is currently effective. + pub fn has(&self, cap: u32) -> bool { + has_capability(self.effective(), cap) + } + + /// Consistent snapshot of the effective bitmask, for callers that gate + /// several capabilities in one decision. + pub fn effective(&self) -> u32 { + self.effective.load(Ordering::SeqCst) + } + + /// Currently-active temporary grants (fresh read of the grants file), + /// for reporting in `SystemInfo`. + pub fn active_grants(&self) -> Vec { + CapabilityGrantStore::load(&self.grants_path).active_grants(now_unix(), self.base) + } + + /// Watch the effective bitmask. Long-running subsystems use this to + /// reconcile in-flight work when a capability appears or disappears. + pub fn subscribe_state(&self) -> watch::Receiver { + self.state_tx.subscribe() + } + + /// Every transition with its change events. Connections forward these to + /// the server as `CapabilitiesChanged`. + pub fn subscribe_transitions(&self) -> broadcast::Receiver { + self.transition_tx.subscribe() + } + + /// Directly set the effective bits, bypassing the grants file. Test-only: + /// lets gate tests exercise a capability flip without a running loop. + #[cfg(test)] + pub fn set_effective_for_test(&self, bits: u32) { + self.effective.store(bits, Ordering::SeqCst); + let _ = self.state_tx.send(bits); + } + + /// Process-wide transition loop: re-reads the grants file every `tick`, + /// updates the effective state, and fans out transitions. Read-only on + /// the file (the CLI is the only writer). Runs for the agent's lifetime. + pub async fn run(self: Arc, tick: Duration) { + // Seed prev_active from the current file so grants already active at + // startup are NOT re-announced as new (avoids alert spam). + let mut prev_active = + CapabilityGrantStore::load(&self.grants_path).active_bits(now_unix(), self.base); + let mut interval = tokio::time::interval(tick); + interval.tick().await; // consume the immediate first tick + + loop { + interval.tick().await; + + let now = now_unix(); + let store = CapabilityGrantStore::load(&self.grants_path); + let (effective, active_bits, temporary, changes) = + evaluate(&store, self.base, prev_active, now); + + if effective != self.effective() { + self.effective.store(effective, Ordering::SeqCst); + let _ = self.state_tx.send(effective); + let _ = self.transition_tx.send(CapabilityTransition { + effective, + temporary, + changes, + }); + tracing::info!(effective, "capability grant state changed"); + } + prev_active = active_bits; + } + } +} + +/// Pure: given the previous active-grant bits and a freshly-loaded store, +/// compute new effective caps, new active bits, the active-grant DTOs, and the +/// change events to emit. +pub fn evaluate( + store: &CapabilityGrantStore, + base: u32, + prev_active_bits: u32, + now: i64, +) -> (u32, u32, Vec, Vec) { + let active_bits = store.active_bits(now, base); + let effective = (base | active_bits) & CAP_VALID_MASK; + let temporary = store.active_grants(now, base); + + let granted = active_bits & !prev_active_bits; + let removed = prev_active_bits & !active_bits; + let mut changes = Vec::new(); + + for meta in ALL_CAPABILITIES { + if granted & meta.bit != 0 { + let rec = store.records().find(|r| r.cap == meta.key); + changes.push(CapabilityChangeEvent { + cap: meta.key.to_string(), + action: CapabilityChangeAction::Granted, + expires_at: rec.map(|r| r.expires_at), + granted_by: rec.map(|r| r.granted_by.clone()), + reason: rec.and_then(|r| r.reason.clone()), + }); + } + if removed & meta.bit != 0 { + // A still-present record means time elapsed (expired); a gone + // record means the operator revoked it. + let rec = store.records().find(|r| r.cap == meta.key); + changes.push(CapabilityChangeEvent { + cap: meta.key.to_string(), + action: if rec.is_some() { + CapabilityChangeAction::Expired + } else { + CapabilityChangeAction::Revoked + }, + expires_at: None, + granted_by: None, + reason: None, + }); + } + } + (effective, active_bits, temporary, changes) +} + +fn now_unix() -> i64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs() as i64) + .unwrap_or(0) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::capability_grants::store::GrantRecord; + use serverbee_common::constants::{CAP_DEFAULT, CAP_TERMINAL}; + + fn store_with(cap: &str, expires_at: i64) -> CapabilityGrantStore { + let mut s = CapabilityGrantStore::default(); + s.upsert( + GrantRecord { + cap: cap.into(), + granted_at: 0, + expires_at, + granted_by: "root".into(), + reason: None, + }, + 0, + ); + s + } + + #[test] + fn newly_active_emits_granted() { + let store = store_with("terminal", 1000); + let (eff, active, temp, changes) = evaluate(&store, CAP_DEFAULT, 0, 0); + assert_eq!(eff, CAP_DEFAULT | CAP_TERMINAL); + assert_eq!(active, CAP_TERMINAL); + assert_eq!(temp.len(), 1); + assert_eq!(temp[0].cap, "terminal"); + assert_eq!(temp[0].expires_at, 1000); + assert_eq!(changes.len(), 1); + assert_eq!(changes[0].action, CapabilityChangeAction::Granted); + assert_eq!(changes[0].cap, "terminal"); + } + + #[test] + fn no_change_when_prev_equals_active() { + let store = store_with("terminal", 1000); + let (_eff, _active, _temp, changes) = evaluate(&store, CAP_DEFAULT, CAP_TERMINAL, 0); + assert!(changes.is_empty()); + } + + #[test] + fn expiry_emits_expired_revoke_emits_revoked() { + let store = store_with("terminal", 100); + let (_e, active, _t, changes) = evaluate(&store, CAP_DEFAULT, CAP_TERMINAL, 200); + assert_eq!(active, 0); + assert_eq!(changes[0].action, CapabilityChangeAction::Expired); + + let empty = CapabilityGrantStore::default(); + let (_e, _a, _t, changes) = evaluate(&empty, CAP_DEFAULT, CAP_TERMINAL, 50); + assert_eq!(changes[0].action, CapabilityChangeAction::Revoked); + } + + #[test] + fn granted_event_carries_grant_metadata() { + // Exercise the `rec.map(...)`/`and_then(...)` arms of the Granted event + // by giving the record a granted_by and a reason. + let mut store = CapabilityGrantStore::default(); + store.upsert( + GrantRecord { + cap: "terminal".into(), + granted_at: 0, + expires_at: 1000, + granted_by: "alice".into(), + reason: Some("incident-42".into()), + }, + 0, + ); + let (_eff, _active, _temp, changes) = evaluate(&store, CAP_DEFAULT, 0, 0); + assert_eq!(changes.len(), 1); + let ev = &changes[0]; + assert_eq!(ev.action, CapabilityChangeAction::Granted); + assert_eq!(ev.expires_at, Some(1000)); + assert_eq!(ev.granted_by.as_deref(), Some("alice")); + assert_eq!(ev.reason.as_deref(), Some("incident-42")); + } + + #[test] + fn evaluate_no_grants_yields_base_effective() { + let store = CapabilityGrantStore::default(); + let (eff, active, temp, changes) = evaluate(&store, CAP_DEFAULT, 0, 0); + assert_eq!(eff, CAP_DEFAULT); + assert_eq!(active, 0); + assert!(temp.is_empty()); + assert!(changes.is_empty()); + } + + #[test] + fn fixed_authority_reflects_base_only() { + let auth = CapabilityAuthority::fixed(CAP_DEFAULT); + assert_eq!(auth.effective(), CAP_DEFAULT); + assert!(!auth.has(CAP_TERMINAL)); + assert!(auth.active_grants().is_empty()); + } + + #[test] + fn new_seeds_effective_from_persisted_grants() { + // A still-active grant in the file is folded into the effective caps + // at construction time, so grants survive an agent restart. + let tmp = tempfile::TempDir::new().unwrap(); + let path = tmp.path().join("capability_grants.json"); + let mut store = CapabilityGrantStore::load(&path); + store.upsert( + GrantRecord { + cap: "terminal".into(), + granted_at: 0, + expires_at: i64::MAX, + granted_by: "root".into(), + reason: None, + }, + 0, + ); + store.flush().unwrap(); + + let auth = CapabilityAuthority::new(CAP_DEFAULT, path); + assert!(auth.has(CAP_TERMINAL)); + assert_eq!(auth.active_grants().len(), 1); + } + + #[tokio::test] + async fn run_fans_out_transition_on_new_grant() { + let tmp = tempfile::TempDir::new().unwrap(); + let path = tmp.path().join("capability_grants.json"); + let store = CapabilityGrantStore::load(&path); + store.flush().unwrap(); + + let auth = CapabilityAuthority::new(CAP_DEFAULT, path.clone()); + let mut transitions = auth.subscribe_transitions(); + let mut state = auth.subscribe_state(); + tokio::spawn(Arc::clone(&auth).run(Duration::from_millis(10))); + + // Let the loop seed `prev_active` from the still-empty file before the + // grant lands, so the transition is observed rather than folded into + // the seed (avoids a race in the Granted event assertion). + tokio::time::sleep(Duration::from_millis(80)).await; + + let mut store = CapabilityGrantStore::load(&path); + store.upsert( + GrantRecord { + cap: "terminal".into(), + granted_at: 0, + expires_at: i64::MAX, + granted_by: "root".into(), + reason: Some("debug".into()), + }, + 0, + ); + store.flush().unwrap(); + + let transition = tokio::time::timeout(Duration::from_secs(3), transitions.recv()) + .await + .expect("transition should arrive within timeout") + .expect("broadcast should deliver"); + assert_eq!(transition.effective, CAP_DEFAULT | CAP_TERMINAL); + assert_eq!(transition.temporary.len(), 1); + assert_eq!(transition.changes.len(), 1); + assert_eq!(transition.changes[0].action, CapabilityChangeAction::Granted); + + // The state watch and the shared snapshot moved with the transition. + tokio::time::timeout(Duration::from_secs(1), state.changed()) + .await + .expect("state watch should flip") + .expect("watch sender must be alive"); + assert_eq!(*state.borrow(), CAP_DEFAULT | CAP_TERMINAL); + assert!(auth.has(CAP_TERMINAL)); + } +} diff --git a/crates/agent/src/capability_grants/mod.rs b/crates/agent/src/capability_grants/mod.rs index 64ef66d96..8b858847e 100644 --- a/crates/agent/src/capability_grants/mod.rs +++ b/crates/agent/src/capability_grants/mod.rs @@ -1,8 +1,8 @@ +pub mod authority; pub mod cli; pub mod store; -pub mod supervisor; -pub use store::CapabilityGrantStore; +pub use authority::CapabilityAuthority; /// Parse a human duration (`90s`, `30m`, `2h`, `1d`) into seconds. Must be a /// positive integer followed by a single unit char. Footgun-guard only. diff --git a/crates/agent/src/capability_grants/supervisor.rs b/crates/agent/src/capability_grants/supervisor.rs deleted file mode 100644 index ab9d89e36..000000000 --- a/crates/agent/src/capability_grants/supervisor.rs +++ /dev/null @@ -1,307 +0,0 @@ -use std::path::PathBuf; -use std::sync::atomic::{AtomicU32, Ordering}; -use std::sync::Arc; -use std::time::Duration; - -use tokio::sync::mpsc; - -use serverbee_common::constants::{ALL_CAPABILITIES, CAP_VALID_MASK}; -use serverbee_common::protocol::{ - AgentMessage, CapabilityChangeAction, CapabilityChangeEvent, TemporaryGrant, -}; - -use super::store::CapabilityGrantStore; - -/// Pure: given the previous active-grant bits and a freshly-loaded store, -/// compute new effective caps, new active bits, the active-grant DTOs, and the -/// change events to emit. -pub fn evaluate( - store: &CapabilityGrantStore, - base: u32, - prev_active_bits: u32, - now: i64, -) -> (u32, u32, Vec, Vec) { - let active_bits = store.active_bits(now, base); - let effective = (base | active_bits) & CAP_VALID_MASK; - let temporary = store.active_grants(now, base); - - let granted = active_bits & !prev_active_bits; - let removed = prev_active_bits & !active_bits; - let mut changes = Vec::new(); - - for meta in ALL_CAPABILITIES { - if granted & meta.bit != 0 { - let rec = store.records().find(|r| r.cap == meta.key); - changes.push(CapabilityChangeEvent { - cap: meta.key.to_string(), - action: CapabilityChangeAction::Granted, - expires_at: rec.map(|r| r.expires_at), - granted_by: rec.map(|r| r.granted_by.clone()), - reason: rec.and_then(|r| r.reason.clone()), - }); - } - if removed & meta.bit != 0 { - // A still-present record means time elapsed (expired); a gone - // record means the operator revoked it. - let rec = store.records().find(|r| r.cap == meta.key); - changes.push(CapabilityChangeEvent { - cap: meta.key.to_string(), - action: if rec.is_some() { - CapabilityChangeAction::Expired - } else { - CapabilityChangeAction::Revoked - }, - expires_at: None, - granted_by: None, - reason: None, - }); - } - } - (effective, active_bits, temporary, changes) -} - -fn now_unix() -> i64 { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_secs() as i64) - .unwrap_or(0) -} - -/// Long-running per-connection task: re-reads the grants file, updates the -/// shared effective-caps cache, and emits `CapabilitiesChanged` on transitions. -/// Read-only on the file (the CLI is the only writer). Stops when `tx` closes -/// (i.e. the connection ended). -pub async fn run_grant_supervisor( - grants_path: PathBuf, - base: u32, - capabilities: Arc, - tx: mpsc::Sender, - tick: Duration, -) { - // Seed prev_active from the current file so grants already active at connect - // time are NOT re-announced as new (avoids alert spam on every reconnect). - let mut prev_active = CapabilityGrantStore::load(&grants_path).active_bits(now_unix(), base); - let mut interval = tokio::time::interval(tick); - interval.tick().await; // consume the immediate first tick - - loop { - tokio::select! { - _ = tx.closed() => { - tracing::debug!("grant supervisor: receiver dropped; stopping"); - break; - } - _ = interval.tick() => {} - } - - let now = now_unix(); - let store = CapabilityGrantStore::load(&grants_path); - let (effective, active_bits, temporary, changes) = - evaluate(&store, base, prev_active, now); - - if effective != capabilities.load(Ordering::SeqCst) { - capabilities.store(effective, Ordering::SeqCst); - let msg = AgentMessage::CapabilitiesChanged { - msg_id: uuid::Uuid::new_v4().to_string(), - capabilities: effective, - temporary, - changes, - }; - if tx.send(msg).await.is_err() { - tracing::debug!("grant supervisor channel closed; stopping"); - break; - } - tracing::info!(effective, "capability grant state changed"); - } - prev_active = active_bits; - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::capability_grants::store::GrantRecord; - use serverbee_common::constants::{CAP_DEFAULT, CAP_TERMINAL}; - - fn store_with(cap: &str, expires_at: i64) -> CapabilityGrantStore { - let mut s = CapabilityGrantStore::default(); - s.upsert( - GrantRecord { - cap: cap.into(), - granted_at: 0, - expires_at, - granted_by: "root".into(), - reason: None, - }, - 0, - ); - s - } - - #[test] - fn newly_active_emits_granted() { - let store = store_with("terminal", 1000); - let (eff, active, temp, changes) = evaluate(&store, CAP_DEFAULT, 0, 0); - assert_eq!(eff, CAP_DEFAULT | CAP_TERMINAL); - assert_eq!(active, CAP_TERMINAL); - assert_eq!(temp.len(), 1); - assert_eq!(temp[0].cap, "terminal"); - assert_eq!(temp[0].expires_at, 1000); - assert_eq!(changes.len(), 1); - assert_eq!(changes[0].action, CapabilityChangeAction::Granted); - assert_eq!(changes[0].cap, "terminal"); - } - - #[test] - fn no_change_when_prev_equals_active() { - let store = store_with("terminal", 1000); - let (_eff, _active, _temp, changes) = evaluate(&store, CAP_DEFAULT, CAP_TERMINAL, 0); - assert!(changes.is_empty()); - } - - #[test] - fn expiry_emits_expired_revoke_emits_revoked() { - let store = store_with("terminal", 100); - let (_e, active, _t, changes) = evaluate(&store, CAP_DEFAULT, CAP_TERMINAL, 200); - assert_eq!(active, 0); - assert_eq!(changes[0].action, CapabilityChangeAction::Expired); - - let empty = CapabilityGrantStore::default(); - let (_e, _a, _t, changes) = evaluate(&empty, CAP_DEFAULT, CAP_TERMINAL, 50); - assert_eq!(changes[0].action, CapabilityChangeAction::Revoked); - } - - #[tokio::test] - async fn supervisor_stops_when_receiver_dropped() { - use serverbee_common::constants::CAP_DEFAULT; - use std::sync::atomic::AtomicU32; - use std::sync::Arc; - - let (tx, rx) = tokio::sync::mpsc::channel(8); - let caps = Arc::new(AtomicU32::new(CAP_DEFAULT)); - let handle = tokio::spawn(run_grant_supervisor( - std::path::PathBuf::from("/nonexistent/capability_grants.json"), - CAP_DEFAULT, - caps, - tx, - std::time::Duration::from_millis(10), - )); - drop(rx); - tokio::time::timeout(std::time::Duration::from_secs(2), handle) - .await - .expect("supervisor should stop promptly after the receiver is dropped") - .expect("supervisor task should not panic"); - } - - #[test] - fn granted_event_carries_grant_metadata() { - // Exercise the `rec.map(...)`/`and_then(...)` arms of the Granted event - // by giving the record a granted_by and a reason. The earlier - // `newly_active_emits_granted` test uses reason: None, leaving the - // Some() branch uncovered. - let mut store = CapabilityGrantStore::default(); - store.upsert( - GrantRecord { - cap: "terminal".into(), - granted_at: 0, - expires_at: 1000, - granted_by: "alice".into(), - reason: Some("incident-42".into()), - }, - 0, - ); - let (_eff, _active, _temp, changes) = evaluate(&store, CAP_DEFAULT, 0, 0); - assert_eq!(changes.len(), 1); - let ev = &changes[0]; - assert_eq!(ev.action, CapabilityChangeAction::Granted); - assert_eq!(ev.expires_at, Some(1000)); - assert_eq!(ev.granted_by.as_deref(), Some("alice")); - assert_eq!(ev.reason.as_deref(), Some("incident-42")); - } - - #[test] - fn evaluate_no_grants_yields_base_effective() { - // Empty store, no previous active bits: effective == base, no changes. - let store = CapabilityGrantStore::default(); - let (eff, active, temp, changes) = evaluate(&store, CAP_DEFAULT, 0, 0); - assert_eq!(eff, CAP_DEFAULT); - assert_eq!(active, 0); - assert!(temp.is_empty()); - assert!(changes.is_empty()); - } - - #[tokio::test] - async fn supervisor_emits_capabilities_changed_on_new_grant() { - use serverbee_common::constants::{CAP_DEFAULT, CAP_TERMINAL}; - use std::sync::atomic::{AtomicU32, Ordering}; - use std::sync::Arc; - - let tmp = tempfile::TempDir::new().unwrap(); - let path = tmp.path().join("capability_grants.json"); - - // Seed an active grant for a cap that is OFF in CAP_DEFAULT (terminal) - // BEFORE starting the supervisor would seed it as prev_active. So instead - // we start with an empty file, then write the grant after start so the - // supervisor observes a transition and emits CapabilitiesChanged. - let store = CapabilityGrantStore::load(&path); - store.flush().unwrap(); - - let (tx, mut rx) = tokio::sync::mpsc::channel(8); - let caps = Arc::new(AtomicU32::new(CAP_DEFAULT)); - let caps_clone = Arc::clone(&caps); - let handle = tokio::spawn(run_grant_supervisor( - path.clone(), - CAP_DEFAULT, - caps_clone, - tx, - std::time::Duration::from_millis(10), - )); - - // Let the supervisor seed `prev_active` from the still-empty file before - // we write the grant; otherwise the grant could land first and be folded - // into the seed, suppressing the Granted change event (test would race). - tokio::time::sleep(std::time::Duration::from_millis(80)).await; - - // Write a grant that turns terminal ON far in the future. - let mut store = CapabilityGrantStore::load(&path); - store.upsert( - GrantRecord { - cap: "terminal".into(), - granted_at: 0, - expires_at: i64::MAX, - granted_by: "root".into(), - reason: Some("debug".into()), - }, - 0, - ); - store.flush().unwrap(); - - let msg = tokio::time::timeout(std::time::Duration::from_secs(3), rx.recv()) - .await - .expect("supervisor should emit within timeout") - .expect("channel should deliver a message"); - - match msg { - AgentMessage::CapabilitiesChanged { - capabilities, - temporary, - changes, - .. - } => { - assert_eq!(capabilities, CAP_DEFAULT | CAP_TERMINAL); - assert_eq!(caps.load(Ordering::SeqCst), CAP_DEFAULT | CAP_TERMINAL); - assert_eq!(temporary.len(), 1); - assert_eq!(temporary[0].cap, "terminal"); - assert_eq!(changes.len(), 1); - assert_eq!(changes[0].action, CapabilityChangeAction::Granted); - } - other => panic!("unexpected message: {other:?}"), - } - - // Drop the receiver so the supervisor stops, then await it. - drop(rx); - tokio::time::timeout(std::time::Duration::from_secs(2), handle) - .await - .expect("supervisor should stop after receiver dropped") - .expect("supervisor task should not panic"); - } -} diff --git a/crates/agent/src/docker/mod.rs b/crates/agent/src/docker/mod.rs index 5faa691c5..449f9739d 100644 --- a/crates/agent/src/docker/mod.rs +++ b/crates/agent/src/docker/mod.rs @@ -6,21 +6,21 @@ pub mod volumes; use std::collections::HashMap; use std::sync::Arc; -use std::sync::atomic::{AtomicU32, Ordering}; use std::time::Duration; use bollard::Docker; -use serverbee_common::constants::{CAP_DOCKER, has_capability}; +use serverbee_common::constants::CAP_DOCKER; use serverbee_common::docker_types::DockerAction; use serverbee_common::protocol::{AgentMessage, ServerMessage}; use tokio::sync::mpsc; use tokio::task::JoinHandle; use tokio::time::Interval; +use crate::capability_grants::CapabilityAuthority; pub struct DockerManager { docker: Docker, agent_tx: mpsc::Sender, - capabilities: Arc, + capabilities: Arc, stats_interval: Option, log_sessions: HashMap>, event_stream_handle: Option>, @@ -32,7 +32,7 @@ impl DockerManager { /// Attempt to connect to the local Docker daemon. pub fn try_new( agent_tx: mpsc::Sender, - capabilities: Arc, + capabilities: Arc, ) -> anyhow::Result { let docker = Docker::connect_with_local_defaults()?; Ok(Self { @@ -87,8 +87,7 @@ impl DockerManager { /// Check whether the Docker capability is currently enabled. fn is_capable(&self) -> bool { - let caps = self.capabilities.load(Ordering::SeqCst); - has_capability(caps, CAP_DOCKER) + self.capabilities.has(CAP_DOCKER) } /// Dispatch a Docker-related server message. @@ -434,7 +433,7 @@ mod tests { let manager = DockerManager { docker, agent_tx: tx, - capabilities: Arc::new(AtomicU32::new(caps)), + capabilities: CapabilityAuthority::fixed(caps), stats_interval: None, log_sessions: HashMap::new(), event_stream_handle: None, @@ -476,7 +475,7 @@ mod tests { // Toggling the shared atomic flips the capability result without rebuilding let (manager, _rx) = make_manager(0); assert!(!manager.is_capable()); - manager.capabilities.store(CAP_DOCKER, Ordering::SeqCst); + manager.capabilities.set_effective_for_test(CAP_DOCKER); assert!(manager.is_capable()); } diff --git a/crates/agent/src/file_manager.rs b/crates/agent/src/file_manager.rs index b36d433d2..627e17e36 100644 --- a/crates/agent/src/file_manager.rs +++ b/crates/agent/src/file_manager.rs @@ -1,6 +1,5 @@ use std::path::PathBuf; use std::sync::Arc; -use std::sync::atomic::AtomicU32; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64; @@ -11,6 +10,7 @@ use serverbee_common::types::{FileEntry, FileType}; use tokio::sync::mpsc; use crate::config::FileConfig; +use crate::capability_grants::CapabilityAuthority; /// Events produced by background file transfer tasks, sent to the reporter loop. #[allow(clippy::enum_variant_names)] @@ -75,13 +75,13 @@ pub struct FileManager { /// Pre-canonicalized root paths, computed once at construction time. canonical_roots: Vec, #[allow(dead_code)] // stored for future per-method capability checks - capabilities: Arc, + capabilities: Arc, active_downloads: DashMap, active_uploads: DashMap, } impl FileManager { - pub fn new(config: FileConfig, capabilities: Arc) -> Self { + pub fn new(config: FileConfig, capabilities: Arc) -> Self { let canonical_roots: Vec = config .root_paths .iter() @@ -733,8 +733,7 @@ async fn download_file( mod tests { use super::*; use serverbee_common::constants::CAP_FILE; - use std::sync::atomic::AtomicU32; - use tempfile::TempDir; + use tempfile::TempDir; fn make_config(root: &str) -> FileConfig { FileConfig { @@ -753,7 +752,7 @@ mod tests { } fn make_manager(config: FileConfig) -> FileManager { - let caps = Arc::new(AtomicU32::new(CAP_FILE)); + let caps = CapabilityAuthority::fixed(CAP_FILE); FileManager::new(config, caps) } @@ -1154,7 +1153,7 @@ mod tests { "passwd".into(), ], }; - let caps = Arc::new(AtomicU32::new(u32::MAX)); + let caps = CapabilityAuthority::fixed(u32::MAX); let mgr = FileManager::new(config, caps); let f1 = tmp1.path().join("a.txt"); diff --git a/crates/agent/src/ip_quality/mod.rs b/crates/agent/src/ip_quality/mod.rs index a930a56cc..26dc9be4f 100644 --- a/crates/agent/src/ip_quality/mod.rs +++ b/crates/agent/src/ip_quality/mod.rs @@ -4,14 +4,15 @@ pub mod rule_engine; pub mod ssrf; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; +use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; use chrono::Utc; -use serverbee_common::constants::{CAP_IP_QUALITY, has_capability}; +use serverbee_common::constants::CAP_IP_QUALITY; use serverbee_common::protocol::{UnlockResultData, UnlockServiceDef, UnlockStatus}; use tokio::sync::{Mutex, Semaphore, mpsc, watch}; use tokio::task::JoinHandle; +use crate::capability_grants::CapabilityAuthority; /// Maximum number of services probed concurrently within a single run. const MAX_UNLOCK_CONCURRENT: usize = 5; @@ -46,7 +47,7 @@ pub struct UnlockChecker { /// Watch channel sender for IP changes (watch: only the latest value matters). ip_changed_tx: watch::Sender>, /// Agent-local capabilities bitmap (shared, immutable at runtime). - capabilities: Arc, + capabilities: Arc, /// Whether the post-connect initial run has been scheduled. The first /// `IpQualitySync` carrying services triggers one immediate run so probes /// don't wait a full `interval_hours` — essential on stable-IP hosts where @@ -61,7 +62,7 @@ impl UnlockChecker { /// /// `result_tx` receives a `RunResult` each time a run completes. pub fn new( - capabilities: Arc, + capabilities: Arc, result_tx: mpsc::Sender, ) -> Self { let services = Arc::new(Mutex::new(Vec::new())); @@ -110,7 +111,7 @@ impl UnlockChecker { /// sends the sync when the capability is effective, but we guard here too /// as defence-in-depth). pub async fn sync(&self, services: Vec, interval_hours: u32) { - if !has_capability(self.capabilities.load(Ordering::SeqCst), CAP_IP_QUALITY) { + if !self.capabilities.has(CAP_IP_QUALITY) { tracing::debug!("IpQualitySync received but CAP_IP_QUALITY not effective — ignoring"); return; } @@ -135,7 +136,7 @@ impl UnlockChecker { /// Trigger an immediate check outside the regular schedule. pub fn run_now(&self) { - if !has_capability(self.capabilities.load(Ordering::SeqCst), CAP_IP_QUALITY) { + if !self.capabilities.has(CAP_IP_QUALITY) { tracing::debug!( "IpQualityRunNow received but CAP_IP_QUALITY not effective — ignoring" ); @@ -318,7 +319,7 @@ async fn run_scheduler( interval_hours: Arc>, mut run_now_rx: mpsc::Receiver<()>, mut ip_changed_rx: watch::Receiver>, - capabilities: Arc, + capabilities: Arc, result_tx: mpsc::Sender, ) { // The last IP value the scheduler has acted on. Updated only by the @@ -371,7 +372,7 @@ async fn run_scheduler( } // Re-check capability before each run. - if !has_capability(capabilities.load(Ordering::SeqCst), CAP_IP_QUALITY) { + if !capabilities.has(CAP_IP_QUALITY) { tracing::debug!("UnlockChecker: CAP_IP_QUALITY not effective, skipping run"); continue; } @@ -406,7 +407,7 @@ async fn run_scheduler( #[cfg(test)] mod tests { use std::sync::Arc; - use std::sync::atomic::{AtomicU32, AtomicUsize}; + use std::sync::atomic::AtomicUsize; use serverbee_common::constants::CAP_IP_QUALITY; use serverbee_common::protocol::{UnlockServiceDef, UnlockStatus}; @@ -414,8 +415,8 @@ mod tests { use super::*; - fn make_caps(bits: u32) -> Arc { - Arc::new(AtomicU32::new(bits)) + fn make_caps(bits: u32) -> Arc { + CapabilityAuthority::fixed(bits) } fn stub_builtin(id: &str, key: &str) -> UnlockServiceDef { diff --git a/crates/agent/src/main.rs b/crates/agent/src/main.rs index 4396b295b..032b42114 100644 --- a/crates/agent/src/main.rs +++ b/crates/agent/src/main.rs @@ -238,26 +238,28 @@ async fn main() -> anyhow::Result<()> { } } - // Start the security pipeline before connecting; it owns a long-lived - // mpsc::Sender that the reporter forwards over the WebSocket. + // Process-wide capability authority: owns the effective bitmask (base + + // temporary grants) and drives every transition. Consumers gate on it; + // its transition loop runs for the agent's lifetime. + let capabilities = crate::capability_grants::CapabilityAuthority::new( + agent_local_capabilities, + config.capabilities.grants_path(), + ); + tokio::spawn(std::sync::Arc::clone(&capabilities).run(std::time::Duration::from_secs(3))); + + // Security pipeline, supervised against the authority so a temporary + // security-events grant can start it and its expiry stops it. It owns a + // long-lived mpsc::Sender that the reporter forwards over the WebSocket. let (security_tx, security_rx) = tokio::sync::mpsc::channel::< serverbee_common::protocol::AgentMessage, >(128); - let _security_manager = match SecurityManager::start( + let _security_supervisor = SecurityManager::spawn_supervised( config.security.clone(), - agent_local_capabilities, + std::sync::Arc::clone(&capabilities), security_tx, - ) - .await - { - Ok(m) => Some(m), - Err(e) => { - tracing::warn!(error = %e, "SecurityManager failed to start; continuing without it"); - None - } - }; + ); - let mut reporter = Reporter::new(config, machine_fingerprint, agent_local_capabilities); + let mut reporter = Reporter::new(config, machine_fingerprint, capabilities); reporter.run_with_external(Some(security_rx)).await; Ok(()) } diff --git a/crates/agent/src/network_prober.rs b/crates/agent/src/network_prober.rs index 4bc2d1be0..dbbd07e01 100644 --- a/crates/agent/src/network_prober.rs +++ b/crates/agent/src/network_prober.rs @@ -1,6 +1,5 @@ use std::collections::HashMap; use std::sync::Arc; -use std::sync::atomic::{AtomicU32, Ordering}; use std::time::Duration; use chrono::Utc; @@ -11,6 +10,8 @@ use tokio::sync::mpsc; use tokio::task::JoinHandle; use tokio::time::MissedTickBehavior; +use crate::capability_grants::CapabilityAuthority; + use crate::probe_utils::{probe_http, probe_icmp_batch, probe_tcp}; /// A running probe task for a single network target. @@ -26,11 +27,14 @@ struct RunningTask { pub struct NetworkProber { tasks: HashMap, tx: mpsc::Sender, - capabilities: Arc, + capabilities: Arc, } impl NetworkProber { - pub fn new(tx: mpsc::Sender, capabilities: Arc) -> Self { + pub fn new( + tx: mpsc::Sender, + capabilities: Arc, + ) -> Self { Self { tasks: HashMap::new(), tx, @@ -45,7 +49,7 @@ impl NetworkProber { /// are (re)started. pub fn sync(&mut self, targets: Vec, interval: u32, packet_count: u32) { // Filter by current capability bitmap - let caps = self.capabilities.load(Ordering::SeqCst); + let caps = self.capabilities.effective(); let targets: Vec<_> = targets .into_iter() .filter(|t| { @@ -354,7 +358,7 @@ mod tests { #[tokio::test] async fn test_network_prober_sync_stops_removed_tasks() { let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(CAP_PING_ICMP | CAP_PING_TCP | CAP_PING_HTTP)); + let caps = CapabilityAuthority::fixed(CAP_PING_ICMP | CAP_PING_TCP | CAP_PING_HTTP); let mut prober = NetworkProber::new(tx, caps); let targets = vec![NetworkProbeTarget { @@ -375,7 +379,7 @@ mod tests { async fn test_network_prober_capability_filter() { let (tx, _rx) = mpsc::channel(16); // Only TCP capability enabled - let caps = Arc::new(AtomicU32::new(CAP_PING_TCP)); + let caps = CapabilityAuthority::fixed(CAP_PING_TCP); let mut prober = NetworkProber::new(tx, caps); let targets = vec![ @@ -546,7 +550,7 @@ mod tests { #[tokio::test] async fn test_sync_restarts_when_interval_changes() { let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(CAP_PING_TCP)); + let caps = CapabilityAuthority::fixed(CAP_PING_TCP); let mut prober = NetworkProber::new(tx, caps); prober.sync(vec![target("t1", "tcp", "1.1.1.1:80")], 60, 3); @@ -565,7 +569,7 @@ mod tests { #[tokio::test] async fn test_sync_restarts_when_packet_count_changes() { let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(CAP_PING_TCP)); + let caps = CapabilityAuthority::fixed(CAP_PING_TCP); let mut prober = NetworkProber::new(tx, caps); prober.sync(vec![target("t1", "tcp", "1.1.1.1:80")], 60, 3); @@ -581,7 +585,7 @@ mod tests { #[tokio::test] async fn test_sync_restarts_when_target_address_changes() { let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(CAP_PING_TCP)); + let caps = CapabilityAuthority::fixed(CAP_PING_TCP); let mut prober = NetworkProber::new(tx, caps); prober.sync(vec![target("t1", "tcp", "1.1.1.1:80")], 60, 3); @@ -598,7 +602,7 @@ mod tests { #[tokio::test] async fn test_sync_restarts_when_probe_type_changes() { let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(CAP_PING_TCP | CAP_PING_HTTP)); + let caps = CapabilityAuthority::fixed(CAP_PING_TCP | CAP_PING_HTTP); let mut prober = NetworkProber::new(tx, caps); prober.sync(vec![target("t1", "tcp", "1.1.1.1:80")], 60, 3); @@ -618,7 +622,7 @@ mod tests { // We can't compare handle identity, but the task count and config stay // stable across the unchanged sync. let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(CAP_PING_TCP)); + let caps = CapabilityAuthority::fixed(CAP_PING_TCP); let mut prober = NetworkProber::new(tx, caps); prober.sync(vec![target("t1", "tcp", "1.1.1.1:80")], 60, 3); @@ -635,7 +639,7 @@ mod tests { #[tokio::test] async fn test_sync_adds_and_keeps_multiple_targets() { let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(CAP_PING_TCP | CAP_PING_HTTP)); + let caps = CapabilityAuthority::fixed(CAP_PING_TCP | CAP_PING_HTTP); let mut prober = NetworkProber::new(tx, caps); prober.sync( @@ -671,7 +675,7 @@ mod tests { async fn test_sync_filters_all_when_no_capabilities() { // Zero capability bitmap -> every target is filtered out. let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(0)); + let caps = CapabilityAuthority::fixed(0); let mut prober = NetworkProber::new(tx, caps); prober.sync( @@ -690,7 +694,7 @@ mod tests { async fn test_sync_filters_unknown_probe_type() { // An unknown probe_type maps to None -> filtered out even with full caps. let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(CAP_PING_ICMP | CAP_PING_TCP | CAP_PING_HTTP)); + let caps = CapabilityAuthority::fixed(CAP_PING_ICMP | CAP_PING_TCP | CAP_PING_HTTP); let mut prober = NetworkProber::new(tx, caps); prober.sync(vec![target("weird", "dns", "example.com")], 60, 3); @@ -700,7 +704,7 @@ mod tests { #[tokio::test] async fn test_stop_all_clears_tasks() { let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(CAP_PING_TCP)); + let caps = CapabilityAuthority::fixed(CAP_PING_TCP); let mut prober = NetworkProber::new(tx, caps); prober.sync(vec![target("t1", "tcp", "1.1.1.1:80")], 60, 3); diff --git a/crates/agent/src/pinger.rs b/crates/agent/src/pinger.rs index 5a90cd0f7..b88b82777 100644 --- a/crates/agent/src/pinger.rs +++ b/crates/agent/src/pinger.rs @@ -1,6 +1,5 @@ use std::collections::HashMap; use std::sync::Arc; -use std::sync::atomic::{AtomicU32, Ordering}; use std::time::Duration; use chrono::Utc; @@ -8,6 +7,7 @@ use serverbee_common::constants::{has_capability, probe_type_to_cap}; use serverbee_common::types::{PingResult, PingTaskConfig}; use tokio::sync::mpsc; +use crate::capability_grants::CapabilityAuthority; use crate::probe_utils; /// Manages running ping probe tasks. Each task runs on its own interval @@ -15,11 +15,14 @@ use crate::probe_utils; pub struct PingManager { tasks: HashMap>, result_tx: mpsc::Sender, - capabilities: Arc, + capabilities: Arc, } impl PingManager { - pub fn new(result_tx: mpsc::Sender, capabilities: Arc) -> Self { + pub fn new( + result_tx: mpsc::Sender, + capabilities: Arc, + ) -> Self { Self { tasks: HashMap::new(), result_tx, @@ -31,7 +34,7 @@ impl PingManager { /// Stops tasks no longer in the list, starts new ones, restarts changed ones. pub fn sync(&mut self, configs: Vec) { // Filter by capability bitmap - let caps = self.capabilities.load(Ordering::SeqCst); + let caps = self.capabilities.effective(); let configs: Vec<_> = configs .into_iter() .filter(|c| { @@ -191,7 +194,7 @@ mod tests { #[tokio::test] async fn test_ping_manager_new_starts_empty() { let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(CAP_PING_ICMP)); + let caps = CapabilityAuthority::fixed(CAP_PING_ICMP); let manager = PingManager::new(tx, caps); assert_eq!(manager.tasks.len(), 0); } @@ -200,7 +203,7 @@ mod tests { async fn test_ping_manager_sync_starts_filtered_tasks() { let (tx, _rx) = mpsc::channel(16); // Only TCP capability enabled. - let caps = Arc::new(AtomicU32::new(CAP_PING_TCP)); + let caps = CapabilityAuthority::fixed(CAP_PING_TCP); let mut manager = PingManager::new(tx, caps); let configs = vec![ @@ -222,7 +225,7 @@ mod tests { #[tokio::test] async fn test_ping_manager_sync_filters_unknown_probe_type() { let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(CAP_PING_ICMP | CAP_PING_TCP | CAP_PING_HTTP)); + let caps = CapabilityAuthority::fixed(CAP_PING_ICMP | CAP_PING_TCP | CAP_PING_HTTP); let mut manager = PingManager::new(tx, caps); // Unknown probe type maps to None -> filtered out. @@ -233,7 +236,7 @@ mod tests { #[tokio::test] async fn test_ping_manager_sync_filters_all_when_no_caps() { let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(0)); + let caps = CapabilityAuthority::fixed(0); let mut manager = PingManager::new(tx, caps); manager.sync(vec![ @@ -246,7 +249,7 @@ mod tests { #[tokio::test] async fn test_ping_manager_sync_stops_removed_tasks() { let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(CAP_PING_TCP)); + let caps = CapabilityAuthority::fixed(CAP_PING_TCP); let mut manager = PingManager::new(tx, caps); manager.sync(vec![config("tcp1", "tcp", "1.1.1.1:80", 30)]); @@ -260,7 +263,7 @@ mod tests { #[tokio::test] async fn test_ping_manager_sync_restarts_running_task() { let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(CAP_PING_TCP)); + let caps = CapabilityAuthority::fixed(CAP_PING_TCP); let mut manager = PingManager::new(tx, caps); manager.sync(vec![config("tcp1", "tcp", "1.1.1.1:80", 30)]); @@ -278,7 +281,7 @@ mod tests { #[tokio::test] async fn test_ping_manager_sync_keeps_unrelated_and_adds_new() { let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(CAP_PING_TCP | CAP_PING_HTTP)); + let caps = CapabilityAuthority::fixed(CAP_PING_TCP | CAP_PING_HTTP); let mut manager = PingManager::new(tx, caps); manager.sync(vec![ @@ -308,7 +311,7 @@ mod tests { // the loop only sends after a tick, we instead drive the finished-handle // branch by aborting via stop_all then re-syncing. let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(CAP_PING_TCP)); + let caps = CapabilityAuthority::fixed(CAP_PING_TCP); let mut manager = PingManager::new(tx, caps); manager.sync(vec![config("tcp1", "tcp", "1.1.1.1:80", 30)]); @@ -333,7 +336,7 @@ mod tests { #[tokio::test] async fn test_ping_manager_stop_all_clears_tasks() { let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(CAP_PING_TCP | CAP_PING_HTTP)); + let caps = CapabilityAuthority::fixed(CAP_PING_TCP | CAP_PING_HTTP); let mut manager = PingManager::new(tx, caps); manager.sync(vec![ @@ -454,7 +457,7 @@ mod tests { // probe_type_to_cap only maps lowercase "icmp"/"tcp"/"http"; an // uppercase variant maps to None and is filtered out even with all caps. let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(CAP_PING_ICMP | CAP_PING_TCP | CAP_PING_HTTP)); + let caps = CapabilityAuthority::fixed(CAP_PING_ICMP | CAP_PING_TCP | CAP_PING_HTTP); let mut manager = PingManager::new(tx, caps); manager.sync(vec![ @@ -469,7 +472,7 @@ mod tests { // Exactly one capability bit set: only the matching probe_type survives, // exercising has_capability against a single-bit mask for every type. let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(CAP_PING_HTTP)); + let caps = CapabilityAuthority::fixed(CAP_PING_HTTP); let mut manager = PingManager::new(tx, caps); manager.sync(vec![ @@ -488,7 +491,7 @@ mod tests { // An empty probe_type string maps to None (catch-all in // probe_type_to_cap) and is dropped regardless of the capability mask. let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(u32::MAX)); + let caps = CapabilityAuthority::fixed(u32::MAX); let mut manager = PingManager::new(tx, caps); manager.sync(vec![config("blank", "", "1.1.1.1", 30)]); @@ -500,7 +503,7 @@ mod tests { // u32::MAX enables every bit, so all three valid probe types pass the // capability filter together. let (tx, _rx) = mpsc::channel(16); - let caps = Arc::new(AtomicU32::new(u32::MAX)); + let caps = CapabilityAuthority::fixed(u32::MAX); let mut manager = PingManager::new(tx, caps); manager.sync(vec![ diff --git a/crates/agent/src/reporter/file_ops.rs b/crates/agent/src/reporter/file_ops.rs index 8052beec8..1c35d4a00 100644 --- a/crates/agent/src/reporter/file_ops.rs +++ b/crates/agent/src/reporter/file_ops.rs @@ -5,16 +5,14 @@ //! one reply-shape rule for denials. `FileDownloadCancel` is exempt: cancel //! must always work so a revoked agent can still be cleaned up. -use std::sync::Arc; -use std::sync::atomic::{AtomicU32, Ordering}; - use futures_util::SinkExt; -use serverbee_common::constants::{CAP_FILE, has_capability}; +use serverbee_common::constants::CAP_FILE; use serverbee_common::protocol::{AgentMessage, ServerMessage}; use tokio::sync::mpsc; use tokio_tungstenite::tungstenite::Message; use super::wire::send_msg; +use crate::capability_grants::CapabilityAuthority; use crate::file_manager::{FileEvent, FileManager}; const DISABLED: &str = "File capability disabled"; @@ -25,7 +23,7 @@ pub(super) async fn handle_file_message( write: &mut S, file_manager: &FileManager, file_tx: &mpsc::Sender, - capabilities: &Arc, + capabilities: &CapabilityAuthority, ) -> anyhow::Result<()> where S: SinkExt + Unpin, @@ -37,8 +35,7 @@ where return Ok(()); } - let caps = capabilities.load(Ordering::SeqCst); - if !has_capability(caps, CAP_FILE) || !file_manager.is_enabled() { + if !capabilities.has(CAP_FILE) || !file_manager.is_enabled() { if let Some(reply) = capability_disabled_reply(msg) { send_msg(write, &reply).await?; } diff --git a/crates/agent/src/reporter/mod.rs b/crates/agent/src/reporter/mod.rs index 2854b55db..07e31b6fd 100644 --- a/crates/agent/src/reporter/mod.rs +++ b/crates/agent/src/reporter/mod.rs @@ -3,7 +3,7 @@ mod wire; use std::net::IpAddr; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; +use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; use futures_util::{SinkExt, StreamExt}; @@ -19,6 +19,7 @@ use tokio_tungstenite::tungstenite::Message; use wire::send_msg; +use crate::capability_grants::CapabilityAuthority; use crate::collector::Collector; use crate::config::AgentConfig; use crate::docker::DockerManager; @@ -41,17 +42,21 @@ static UPGRADE_IN_PROGRESS: AtomicBool = AtomicBool::new(false); pub struct Reporter { config: AgentConfig, fingerprint: String, - agent_local_capabilities: u32, + capabilities: Arc, firewall_manager: Arc, } impl Reporter { - pub fn new(config: AgentConfig, fingerprint: String, agent_local_capabilities: u32) -> Self { + pub fn new( + config: AgentConfig, + fingerprint: String, + capabilities: Arc, + ) -> Self { let firewall_manager = Arc::new(FirewallManager::new(Arc::new(CliNftExecutor))); Self { config, fingerprint, - agent_local_capabilities, + capabilities, firewall_manager, } } @@ -137,19 +142,12 @@ impl Reporter { tracing::info!("Connecting to {}...", build_ws_url(&self.config)?); - // Fold any active temporary grants into the initial effective caps so the - // first SystemInfo already reflects grants that survived a restart. - let base_caps = self.agent_local_capabilities; - let grants_path = self.config.capabilities.grants_path(); - let now0 = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_secs() as i64) - .unwrap_or(0); - let grant_store0 = crate::capability_grants::CapabilityGrantStore::load(&grants_path); - let effective_caps = - (base_caps | grant_store0.active_bits(now0, base_caps)) & serverbee_common::constants::CAP_VALID_MASK; - let initial_temporary = grant_store0.active_grants(now0, base_caps); - let capabilities = Arc::new(AtomicU32::new(effective_caps)); + // The process-wide authority already folds active temporary grants + // into the effective caps, so the first SystemInfo reflects grants + // that survived a restart. + let capabilities = Arc::clone(&self.capabilities); + let effective_caps = capabilities.effective(); + let initial_temporary = capabilities.active_grants(); let request = build_ws_request(&self.config)?; let (ws_stream, _response) = connect_async(request).await?; @@ -283,7 +281,7 @@ impl Reporter { // arm of the select! loop below. { use serverbee_common::constants::CAP_IP_QUALITY; - if has_capability(capabilities.load(Ordering::SeqCst), CAP_IP_QUALITY) { + if capabilities.has(CAP_IP_QUALITY) { let seed_ip = initial_ipv4.clone().or_else(|| initial_ipv6.clone()); // Only seed when we have a non-empty interface-derived IP. // If neither is available the checker stays at None and will be @@ -301,24 +299,10 @@ impl Reporter { // Channel for background command execution results. let (cmd_result_tx, mut cmd_result_rx) = mpsc::channel::(32); - // Capability grant supervisor: re-reads the grants file and pushes - // CapabilitiesChanged through `grant_tx` (forwarded onto the WS below). - let (grant_tx, mut grant_rx) = mpsc::channel::(8); - { - let grants_path = grants_path.clone(); - let caps = Arc::clone(&capabilities); - let tx = grant_tx.clone(); - tokio::spawn(async move { - crate::capability_grants::supervisor::run_grant_supervisor( - grants_path, - base_caps, - caps, - tx, - std::time::Duration::from_secs(3), - ) - .await; - }); - } + // Capability transitions from the process-wide authority: each one is + // forwarded to the server as CapabilitiesChanged, after reconciling + // in-flight work owned by this connection. + let mut transition_rx = capabilities.subscribe_transitions(); // Fire-and-forget: refine the just-sent SystemInfo IPs with an // externally-observed public IP. If discovery yields something @@ -372,7 +356,7 @@ impl Reporter { // Intercept IpChanged to notify the UnlockChecker before forwarding. if let AgentMessage::IpChanged { ref ipv4, ref ipv6, .. } = cmd_msg { use serverbee_common::constants::CAP_IP_QUALITY; - if has_capability(capabilities.load(Ordering::SeqCst), CAP_IP_QUALITY) { + if capabilities.has(CAP_IP_QUALITY) { // Prefer IPv4 egress; fall back to IPv6. let new_ip = ipv4.clone().or_else(|| ipv6.clone()); unlock_checker.notify_ip_changed(new_ip); @@ -381,9 +365,49 @@ impl Reporter { send_msg(&mut write, &cmd_msg).await?; tracing::debug!("Sent background command result"); } - Some(grant_msg) = grant_rx.recv() => { - send_msg(&mut write, &grant_msg).await?; - tracing::debug!("Sent CapabilitiesChanged"); + transition = transition_rx.recv() => { + match transition { + Ok(t) => { + // Reconcile in-flight work before announcing: a + // revoked or expired terminal capability must tear + // down live PTY sessions, not just block new ones. + if !has_capability(t.effective, CAP_TERMINAL) { + for session_id in terminal_manager.close_all() { + let msg = AgentMessage::TerminalError { + session_id, + error: "Terminal capability revoked or expired".to_string(), + }; + send_msg(&mut write, &msg).await?; + } + } + let msg = AgentMessage::CapabilitiesChanged { + msg_id: uuid::Uuid::new_v4().to_string(), + capabilities: t.effective, + temporary: t.temporary, + changes: t.changes, + }; + send_msg(&mut write, &msg).await?; + tracing::debug!("Sent CapabilitiesChanged"); + } + Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => { + // Missed transitions: resync the server with the + // current snapshot (change events for the missed + // steps are lost, but state converges). + tracing::warn!("capability transitions lagged by {n}; resyncing"); + let msg = AgentMessage::CapabilitiesChanged { + msg_id: uuid::Uuid::new_v4().to_string(), + capabilities: capabilities.effective(), + temporary: capabilities.active_grants(), + changes: Vec::new(), + }; + send_msg(&mut write, &msg).await?; + } + Err(tokio::sync::broadcast::error::RecvError::Closed) => { + // The authority lives as long as the process; a + // closed channel means shutdown is underway. + return Ok(()); + } + } } Some(run_result) = unlock_result_rx.recv() => { let msg = AgentMessage::UnlockResults { @@ -574,7 +598,7 @@ impl Reporter { terminal_manager: &mut TerminalManager, network_prober: &mut NetworkProber, cmd_result_tx: &mpsc::Sender, - capabilities: &Arc, + capabilities: &CapabilityAuthority, file_manager: &FileManager, file_tx: &mpsc::Sender, docker_manager: &mut Option, @@ -605,8 +629,7 @@ impl Reporter { command, timeout, } => { - let caps = capabilities.load(Ordering::SeqCst); - if !has_capability(caps, CAP_EXEC) { + if !capabilities.has(CAP_EXEC) { tracing::warn!("Exec denied: capability disabled (task_id={task_id})"); spawn_capability_denied(cmd_result_tx, Some(task_id), "exec"); return Ok(()); @@ -666,8 +689,7 @@ impl Reporter { job_id, .. } => { - let caps = capabilities.load(Ordering::SeqCst); - if !has_capability(caps, CAP_UPGRADE) { + if !capabilities.has(CAP_UPGRADE) { tracing::warn!("Upgrade denied: capability disabled"); send_msg(write, &capability_denied_msg(None, "upgrade")).await?; return Ok(()); @@ -723,8 +745,7 @@ impl Reporter { max_hops, protocol, } => { - let caps = capabilities.load(Ordering::SeqCst); - if !has_capability(caps, CAP_PING_ICMP) { + if !capabilities.has(CAP_PING_ICMP) { tracing::warn!( "Traceroute denied: capability disabled (request_id={request_id})" ); @@ -849,8 +870,7 @@ impl Reporter { | ServerMessage::BlocklistRemove { .. } | ServerMessage::BlocklistReset) => { let is_reset = matches!(msg, ServerMessage::BlocklistReset); - let caps = capabilities.load(Ordering::SeqCst); - if !is_reset && !has_capability(caps, CAP_FIREWALL_BLOCK) { + if !is_reset && !capabilities.has(CAP_FIREWALL_BLOCK) { tracing::warn!( "Firewall blocklist mutation denied: CAP_FIREWALL_BLOCK not effective — ignoring" ); @@ -860,8 +880,7 @@ impl Reporter { } } ServerMessage::IpQualitySync { services, interval_hours } => { - let caps = capabilities.load(Ordering::SeqCst); - if has_capability(caps, CAP_IP_QUALITY) { + if capabilities.has(CAP_IP_QUALITY) { tracing::info!( "Received IpQualitySync: {} services, interval={}h", services.len(), @@ -875,8 +894,7 @@ impl Reporter { } } ServerMessage::IpQualityRunNow => { - let caps = capabilities.load(Ordering::SeqCst); - if has_capability(caps, CAP_IP_QUALITY) { + if capabilities.has(CAP_IP_QUALITY) { tracing::info!("Received IpQualityRunNow"); unlock_checker.run_now(); } else { @@ -1650,7 +1668,7 @@ mod tests { /// plus the receiver ends so background senders never see a closed channel. struct Harness { reporter: Reporter, - capabilities: Arc, + capabilities: Arc, ping_manager: PingManager, terminal_manager: TerminalManager, network_prober: NetworkProber, @@ -1686,8 +1704,8 @@ mod tests { security: SecurityConfig::default(), capabilities: CapabilitiesConfig::default(), }; - let reporter = Reporter::new(config, "fp".to_string(), caps); - let capabilities = Arc::new(AtomicU32::new(caps)); + let capabilities = CapabilityAuthority::fixed(caps); + let reporter = Reporter::new(config, "fp".to_string(), Arc::clone(&capabilities)); let (ping_tx, _ping_rx) = mpsc::channel(16); let ping_manager = PingManager::new(ping_tx, Arc::clone(&capabilities)); @@ -2810,7 +2828,7 @@ mod tests { // returns Ok. let (listener, addr) = bind_fake_server().await; let tmp = tempfile::tempdir().unwrap(); - let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), ALL_CAPS); + let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), CapabilityAuthority::fixed(ALL_CAPS)); let server = tokio::spawn(async move { let mut ws = accept_ws(&listener).await; @@ -2855,7 +2873,7 @@ mod tests { // the fake server reads off the wire. let (listener, addr) = bind_fake_server().await; let tmp = tempfile::tempdir().unwrap(); - let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), ALL_CAPS); + let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), CapabilityAuthority::fixed(ALL_CAPS)); let server = tokio::spawn(async move { let mut ws = accept_ws(&listener).await; @@ -2890,7 +2908,7 @@ mod tests { // agent Pong over the real socket. let (listener, addr) = bind_fake_server().await; let tmp = tempfile::tempdir().unwrap(); - let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), ALL_CAPS); + let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), CapabilityAuthority::fixed(ALL_CAPS)); let server = tokio::spawn(async move { let mut ws = accept_ws(&listener).await; @@ -2919,7 +2937,7 @@ mod tests { // select! loop. let (listener, addr) = bind_fake_server().await; let tmp = tempfile::tempdir().unwrap(); - let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), ALL_CAPS); + let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), CapabilityAuthority::fixed(ALL_CAPS)); let server = tokio::spawn(async move { let mut ws = accept_ws(&listener).await; @@ -2968,7 +2986,7 @@ mod tests { // receive-dispatch path for a "silent" variant plus the Close arm. let (listener, addr) = bind_fake_server().await; let tmp = tempfile::tempdir().unwrap(); - let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), ALL_CAPS); + let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), CapabilityAuthority::fixed(ALL_CAPS)); let server = tokio::spawn(async move { let mut ws = accept_ws(&listener).await; @@ -2999,7 +3017,7 @@ mod tests { // Ok(()) (the normal-reconnect signal), not an error. let (listener, addr) = bind_fake_server().await; let tmp = tempfile::tempdir().unwrap(); - let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), ALL_CAPS); + let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), CapabilityAuthority::fixed(ALL_CAPS)); let server = tokio::spawn(async move { let mut ws = accept_ws(&listener).await; @@ -3033,7 +3051,7 @@ mod tests { // Drop the listener so the port refuses connections. drop(listener); let tmp = tempfile::tempdir().unwrap(); - let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), ALL_CAPS); + let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), CapabilityAuthority::fixed(ALL_CAPS)); let connect = run_connect_once(&mut reporter, Duration::from_secs(10)).await; let connect = connect.expect("connect should fail fast, not hang"); @@ -3052,7 +3070,7 @@ mod tests { // is bounded by an outer timeout that aborts the never-returning task. let (listener, addr) = bind_fake_server().await; let tmp = tempfile::tempdir().unwrap(); - let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), ALL_CAPS); + let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), CapabilityAuthority::fixed(ALL_CAPS)); // Server: accept two connections; each time send Welcome, read // SystemInfo, then close. Signal each successful handshake. @@ -3131,7 +3149,7 @@ mod tests { // the loop. let (listener, addr) = bind_fake_server().await; let tmp = tempfile::tempdir().unwrap(); - let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), ALL_CAPS); + let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), CapabilityAuthority::fixed(ALL_CAPS)); let server = tokio::spawn(async move { let mut ws = accept_ws(&listener).await; @@ -3164,7 +3182,7 @@ mod tests { // No WS frame results; we confirm liveness with Ping->Pong. let (listener, addr) = bind_fake_server().await; let tmp = tempfile::tempdir().unwrap(); - let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), ALL_CAPS); + let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), CapabilityAuthority::fixed(ALL_CAPS)); let server = tokio::spawn(async move { let mut ws = accept_ws(&listener).await; @@ -3202,7 +3220,7 @@ mod tests { let (listener, addr) = bind_fake_server().await; let tmp = tempfile::tempdir().unwrap(); let caps = ALL_CAPS & !serverbee_common::constants::CAP_TERMINAL; - let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), caps); + let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), CapabilityAuthority::fixed(caps)); let server = tokio::spawn(async move { let mut ws = accept_ws(&listener).await; @@ -3261,6 +3279,110 @@ mod tests { assert!(matches!(pong, AgentMessage::Pong)); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_e2e_terminal_grant_expiry_tears_down_live_session() { + // A live PTY session opened under a temporary terminal grant must be + // torn down when the grant expires: the reconcile step in the + // transition arm closes the session (TerminalError over the WS) and + // the CapabilitiesChanged announcement drops the terminal bit. + use crate::capability_grants::store::GrantRecord; + + let (listener, addr) = bind_fake_server().await; + let tmp = tempfile::tempdir().unwrap(); + let base = ALL_CAPS & !serverbee_common::constants::CAP_TERMINAL; + + // Seed a grant that expires shortly after the session is opened. + let config = e2e_config(&addr, tmp.path()); + let grants_path = config.capabilities.grants_path(); + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs() as i64; + let mut store = crate::capability_grants::store::CapabilityGrantStore::load(&grants_path); + store.upsert( + GrantRecord { + cap: "terminal".into(), + granted_at: now, + // Generous window: the PTY must be observably started before + // the grant expires, even on a loaded CI machine. + expires_at: now + 5, + granted_by: "root".into(), + reason: Some("e2e".into()), + }, + now, + ); + store.flush().unwrap(); + + let authority = CapabilityAuthority::new(base, grants_path); + tokio::spawn(Arc::clone(&authority).run(Duration::from_millis(100))); + let mut reporter = Reporter::new(config, "fp".to_string(), authority); + + let server = tokio::spawn(async move { + let mut ws = accept_ws(&listener).await; + send_welcome(&mut ws, 30).await; + let info = handshake_collect_system_info(&mut ws).await; + send_server_msg( + &mut ws, + &ServerMessage::TerminalOpen { + session_id: "term-live".to_string(), + rows: 24, + cols: 80, + }, + ) + .await; + // The grant is active, so the PTY actually starts. + let _started = read_agent_until(&mut ws, |m| { + matches!(m, AgentMessage::TerminalStarted { session_id } if session_id == "term-live") + }) + .await; + // On expiry the reconcile step closes the session... + let err = read_agent_until(&mut ws, |m| { + matches!( + m, + AgentMessage::TerminalError { session_id, error } + if session_id == "term-live" && error.contains("capability") + ) + }) + .await; + // ...and the transition is announced without the terminal bit. + let changed = read_agent_until(&mut ws, |m| { + matches!(m, AgentMessage::CapabilitiesChanged { .. }) + }) + .await; + ws.send(WsMessage::Close(None)).await.ok(); + (info, err, changed) + }); + + let (info, _err, changed) = drive_e2e(&mut reporter, server, Duration::from_secs(15)).await; + + // The initial SystemInfo carried the granted terminal capability. + match info { + AgentMessage::SystemInfo { + agent_local_capabilities, + .. + } => assert_eq!( + agent_local_capabilities, + Some(ALL_CAPS), + "grant must be folded into the caps reported at connect" + ), + other => panic!("expected SystemInfo, got {other:?}"), + } + match changed { + AgentMessage::CapabilitiesChanged { + capabilities, + changes, + .. + } => { + assert_eq!(capabilities, base, "expiry must drop the terminal bit"); + assert!( + changes.iter().any(|c| c.cap == "terminal"), + "the transition must carry the terminal change event" + ); + } + other => panic!("expected CapabilitiesChanged, got {other:?}"), + } + } + #[tokio::test] async fn test_e2e_dispatch_blocklist_reset_forwards_ack() { // BlocklistReset routes into the FirewallManager and forwards its @@ -3269,7 +3391,7 @@ mod tests { // ack, so the dispatcher always emits exactly one reply frame. let (listener, addr) = bind_fake_server().await; let tmp = tempfile::tempdir().unwrap(); - let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), ALL_CAPS); + let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), CapabilityAuthority::fixed(ALL_CAPS)); let server = tokio::spawn(async move { let mut ws = accept_ws(&listener).await; @@ -3295,7 +3417,7 @@ mod tests { // `nft` but the manager still returns a (failed-state) ack frame. let (listener, addr) = bind_fake_server().await; let tmp = tempfile::tempdir().unwrap(); - let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), ALL_CAPS); + let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), CapabilityAuthority::fixed(ALL_CAPS)); let server = tokio::spawn(async move { let mut ws = accept_ws(&listener).await; @@ -3325,7 +3447,7 @@ mod tests { let (listener, addr) = bind_fake_server().await; let tmp = tempfile::tempdir().unwrap(); let caps = ALL_CAPS & !serverbee_common::constants::CAP_FILE; - let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), caps); + let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), CapabilityAuthority::fixed(caps)); let server = tokio::spawn(async move { let mut ws = accept_ws(&listener).await; @@ -3371,7 +3493,7 @@ mod tests { let state = tempfile::tempdir().unwrap(); let mut config = e2e_config(&addr, state.path()); config.file = enabled_file_cfg(&root); - let mut reporter = Reporter::new(config, "fp".to_string(), ALL_CAPS); + let mut reporter = Reporter::new(config, "fp".to_string(), CapabilityAuthority::fixed(ALL_CAPS)); let dest = root.join("upload.bin"); let dest_s = dest.to_string_lossy().to_string(); @@ -3413,7 +3535,7 @@ mod tests { let (listener, addr) = bind_fake_server().await; let tmp = tempfile::tempdir().unwrap(); let caps = ALL_CAPS & !serverbee_common::constants::CAP_EXEC; - let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), caps); + let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), CapabilityAuthority::fixed(caps)); let server = tokio::spawn(async move { let mut ws = accept_ws(&listener).await; @@ -3459,7 +3581,7 @@ mod tests { let (listener, addr) = bind_fake_server().await; let tmp = tempfile::tempdir().unwrap(); let caps = ALL_CAPS & !serverbee_common::constants::CAP_UPGRADE; - let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), caps); + let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), CapabilityAuthority::fixed(caps)); let server = tokio::spawn(async move { let mut ws = accept_ws(&listener).await; @@ -3502,7 +3624,7 @@ mod tests { // channel and forwarded over the WS. No traceroute subprocess spawns. let (listener, addr) = bind_fake_server().await; let tmp = tempfile::tempdir().unwrap(); - let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), ALL_CAPS); + let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), CapabilityAuthority::fixed(ALL_CAPS)); let server = tokio::spawn(async move { let mut ws = accept_ws(&listener).await; @@ -3544,7 +3666,7 @@ mod tests { // running, proven by a subsequent Ping->Pong. let (listener, addr) = bind_fake_server().await; let tmp = tempfile::tempdir().unwrap(); - let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), ALL_CAPS); + let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), CapabilityAuthority::fixed(ALL_CAPS)); let server = tokio::spawn(async move { let mut ws = accept_ws(&listener).await; @@ -3572,7 +3694,7 @@ mod tests { // dispatching application messages via a ServerMessage::Ping->Pong. let (listener, addr) = bind_fake_server().await; let tmp = tempfile::tempdir().unwrap(); - let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), ALL_CAPS); + let mut reporter = Reporter::new(e2e_config(&addr, tmp.path()), "fp".to_string(), CapabilityAuthority::fixed(ALL_CAPS)); let server = tokio::spawn(async move { let mut ws = accept_ws(&listener).await; diff --git a/crates/agent/src/security/manager.rs b/crates/agent/src/security/manager.rs index 20bfbb3cd..2e3a4b9c5 100644 --- a/crates/agent/src/security/manager.rs +++ b/crates/agent/src/security/manager.rs @@ -16,6 +16,7 @@ use tokio::sync::Mutex; use tokio::sync::mpsc; use tokio::task::JoinHandle; +use crate::capability_grants::CapabilityAuthority; use crate::config::SecurityConfig; use crate::security::conntrack_watcher::{self, ConntrackEvent}; use crate::security::first_seen_store::FirstSeenStore; @@ -130,6 +131,45 @@ impl SecurityManager { pub fn handle_count(&self) -> usize { self.handles.len() } + + /// Supervise the pipeline against the capability authority. + /// + /// Starts the pipeline when `CAP_SECURITY_EVENTS` is (or becomes) + /// effective — including via a temporary grant at any point in the + /// agent's lifetime — and stops it when the capability expires or is + /// revoked, so a bounded grant tears its running work down. Replaces the + /// old start-once-at-boot wiring, which could neither start on a later + /// grant nor stop on expiry. + pub fn spawn_supervised( + cfg: SecurityConfig, + authority: Arc, + tx: mpsc::Sender, + ) -> JoinHandle<()> { + tokio::spawn(async move { + let mut state = authority.subscribe_state(); + let mut running: Option = None; + loop { + let wanted = has_capability(*state.borrow_and_update(), CAP_SECURITY_EVENTS); + if wanted && running.is_none() { + match Self::start(cfg.clone(), authority.effective(), tx.clone()).await { + Ok(m) => running = Some(m), + Err(e) => tracing::warn!( + error = %e, + "SecurityManager failed to start; will retry on next capability change" + ), + } + } else if !wanted && running.take().is_some() { + // Dropping the manager aborts its pipeline tasks. + tracing::info!( + "security pipeline stopped: CAP_SECURITY_EVENTS no longer effective" + ); + } + if state.changed().await.is_err() { + break; + } + } + }) + } } impl Drop for SecurityManager { diff --git a/crates/agent/src/terminal.rs b/crates/agent/src/terminal.rs index 57a00ea46..e576885e3 100644 --- a/crates/agent/src/terminal.rs +++ b/crates/agent/src/terminal.rs @@ -1,14 +1,15 @@ use std::collections::HashMap; use std::io::{Read, Write}; use std::sync::Arc; -use std::sync::atomic::{AtomicU32, Ordering}; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64; use portable_pty::{CommandBuilder, NativePtySystem, PtySize, PtySystem}; -use serverbee_common::constants::{CAP_TERMINAL, MAX_TERMINAL_SESSIONS, has_capability}; +use serverbee_common::constants::{CAP_TERMINAL, MAX_TERMINAL_SESSIONS}; use tokio::sync::mpsc; +use crate::capability_grants::CapabilityAuthority; + /// Message sent from terminal sessions back to the reporter. pub enum TerminalEvent { Output { session_id: String, data: String }, @@ -28,11 +29,14 @@ struct PtySession { pub struct TerminalManager { sessions: HashMap, event_tx: mpsc::Sender, - capabilities: Arc, + capabilities: Arc, } impl TerminalManager { - pub fn new(event_tx: mpsc::Sender, capabilities: Arc) -> Self { + pub fn new( + event_tx: mpsc::Sender, + capabilities: Arc, + ) -> Self { Self { sessions: HashMap::new(), event_tx, @@ -42,8 +46,7 @@ impl TerminalManager { /// Open a new terminal session with the given dimensions. pub fn open(&mut self, session_id: String, rows: u16, cols: u16) { - let caps = self.capabilities.load(Ordering::SeqCst); - if !has_capability(caps, CAP_TERMINAL) { + if !self.capabilities.has(CAP_TERMINAL) { tracing::warn!("Terminal denied: capability disabled (session={session_id})"); let tx = self.event_tx.clone(); let sid = session_id; @@ -340,7 +343,7 @@ mod tests { async fn open_denied_without_capability_emits_error() { let (tx, mut rx) = mpsc::channel(64); // No CAP_TERMINAL bit set. - let caps = Arc::new(AtomicU32::new(0)); + let caps = CapabilityAuthority::fixed(0); let mut mgr = TerminalManager::new(tx, caps); let sid = "denied".to_string(); @@ -375,7 +378,7 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn open_emits_started_and_tracks_session() { let (tx, mut rx) = mpsc::channel(64); - let caps = Arc::new(AtomicU32::new(CAP_TERMINAL)); + let caps = CapabilityAuthority::fixed(CAP_TERMINAL); let mut mgr = TerminalManager::new(tx, caps); let sid = "started".to_string(); @@ -406,7 +409,7 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn open_duplicate_session_id_is_noop() { let (tx, mut rx) = mpsc::channel(64); - let caps = Arc::new(AtomicU32::new(CAP_TERMINAL)); + let caps = CapabilityAuthority::fixed(CAP_TERMINAL); let mut mgr = TerminalManager::new(tx, caps); let sid = "dup".to_string(); @@ -449,7 +452,7 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn open_beyond_max_sessions_emits_error() { let (tx, mut rx) = mpsc::channel(128); - let caps = Arc::new(AtomicU32::new(CAP_TERMINAL)); + let caps = CapabilityAuthority::fixed(CAP_TERMINAL); let mut mgr = TerminalManager::new(tx, caps); // Fill up to the cap. @@ -504,7 +507,7 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn write_input_valid_keeps_session() { let (tx, mut rx) = mpsc::channel(64); - let caps = Arc::new(AtomicU32::new(CAP_TERMINAL)); + let caps = CapabilityAuthority::fixed(CAP_TERMINAL); let mut mgr = TerminalManager::new(tx, caps); let sid = "write".to_string(); @@ -533,7 +536,7 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn write_input_invalid_base64_is_noop() { let (tx, mut rx) = mpsc::channel(64); - let caps = Arc::new(AtomicU32::new(CAP_TERMINAL)); + let caps = CapabilityAuthority::fixed(CAP_TERMINAL); let mut mgr = TerminalManager::new(tx, caps); let sid = "bad-b64".to_string(); @@ -560,7 +563,7 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn write_input_unknown_session_is_noop() { let (tx, _rx) = mpsc::channel(64); - let caps = Arc::new(AtomicU32::new(CAP_TERMINAL)); + let caps = CapabilityAuthority::fixed(CAP_TERMINAL); let mut mgr = TerminalManager::new(tx, caps); // No session opened: hits the `None` branch of get_mut. @@ -572,7 +575,7 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn resize_existing_session_succeeds() { let (tx, mut rx) = mpsc::channel(64); - let caps = Arc::new(AtomicU32::new(CAP_TERMINAL)); + let caps = CapabilityAuthority::fixed(CAP_TERMINAL); let mut mgr = TerminalManager::new(tx, caps); let sid = "resize".to_string(); @@ -597,7 +600,7 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn resize_unknown_session_is_noop() { let (tx, _rx) = mpsc::channel(64); - let caps = Arc::new(AtomicU32::new(CAP_TERMINAL)); + let caps = CapabilityAuthority::fixed(CAP_TERMINAL); let mut mgr = TerminalManager::new(tx, caps); // Hits the `None` branch of get_mut in resize(). @@ -609,7 +612,7 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn close_unknown_session_is_noop() { let (tx, _rx) = mpsc::channel(64); - let caps = Arc::new(AtomicU32::new(CAP_TERMINAL)); + let caps = CapabilityAuthority::fixed(CAP_TERMINAL); let mut mgr = TerminalManager::new(tx, caps); // `sessions.remove` returns None -> the `if let Some` body is skipped. @@ -622,7 +625,7 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn close_all_returns_ids_and_clears_sessions() { let (tx, mut rx) = mpsc::channel(128); - let caps = Arc::new(AtomicU32::new(CAP_TERMINAL)); + let caps = CapabilityAuthority::fixed(CAP_TERMINAL); let mut mgr = TerminalManager::new(tx, caps); let ids = ["a", "b"]; @@ -672,7 +675,7 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn close_all_empty_returns_empty() { let (tx, _rx) = mpsc::channel(64); - let caps = Arc::new(AtomicU32::new(CAP_TERMINAL)); + let caps = CapabilityAuthority::fixed(CAP_TERMINAL); let mut mgr = TerminalManager::new(tx, caps); assert!(mgr.close_all().is_empty()); @@ -681,7 +684,7 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn close_reaps_child_process() { let (tx, _rx) = mpsc::channel(64); - let caps = Arc::new(AtomicU32::new(CAP_TERMINAL)); + let caps = CapabilityAuthority::fixed(CAP_TERMINAL); let mut mgr = TerminalManager::new(tx, caps); let sid = "reap-test".to_string(); @@ -730,7 +733,7 @@ mod tests { async fn open_denied_with_other_capability_only() { let (tx, mut rx) = mpsc::channel(64); // A different bit is on (CAP_EXEC), but CAP_TERMINAL is absent. - let caps = Arc::new(AtomicU32::new(CAP_EXEC)); + let caps = CapabilityAuthority::fixed(CAP_EXEC); let mut mgr = TerminalManager::new(tx, caps); let sid = "other-cap".to_string(); @@ -759,17 +762,17 @@ mod tests { } } - /// The capability is re-read from the shared atomic on every `open()` call: + /// The capability is re-read from the authority on every `open()` call: /// revoking CAP_TERMINAL after construction (e.g. capability revocation) /// must deny a subsequent open without spawning a PTY. #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn open_respects_runtime_capability_revocation() { let (tx, mut rx) = mpsc::channel(64); - let caps = Arc::new(AtomicU32::new(CAP_TERMINAL)); + let caps = CapabilityAuthority::fixed(CAP_TERMINAL); let mut mgr = TerminalManager::new(tx, Arc::clone(&caps)); - // Revoke the terminal capability through the shared handle before opening. - caps.store(0, Ordering::SeqCst); + // Revoke the terminal capability through the shared authority before opening. + caps.set_effective_for_test(0); let sid = "revoked".to_string(); mgr.open(sid.clone(), 24, 80); @@ -800,7 +803,7 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn write_input_empty_base64_unknown_session_is_noop() { let (tx, _rx) = mpsc::channel(64); - let caps = Arc::new(AtomicU32::new(CAP_TERMINAL)); + let caps = CapabilityAuthority::fixed(CAP_TERMINAL); let mut mgr = TerminalManager::new(tx, caps); // "" is valid base64 for an empty byte slice; session does not exist. From c4c53da2c7a366737373d82a1868cd8bde9721d5 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 10:45:16 +0800 Subject: [PATCH 04/44] docs: record monitor check and capability lifecycle fixes Add Unreleased changelog entries for the shared WebSocket credential policy, the unified service-monitor check transition, and temporary grants bounding running work. Document the new expiry semantics (terminal teardown, live security-pipeline start/stop) on the capabilities page in both languages. --- CHANGELOG.md | 10 ++++++++++ apps/docs/content/docs/en/capabilities.mdx | 1 + apps/docs/content/docs/zh/capabilities.mdx | 1 + 3 files changed, 12 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index aba7ed884..4672dfe94 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Security + +- **WebSocket handlers share the HTTP credential policy** -- The browser, terminal, and Docker-logs WebSocket endpoints each carried a private copy of the credential ladder (cookie → API key → Bearer), so precedence, forced-password, and mobile-expiry rules could silently drift between four copies. All transports now resolve through one shared policy in the auth middleware. Two edge cases tightened along the way: a user flagged for a forced password change can no longer fall through from their session cookie to a weaker credential on WebSocket upgrades, and a fixed-lifetime mobile session now enforces its expiry on long-lived sockets whichever header carried it + +### Fixed + +- **Manual service-monitor checks follow the full check policy** -- The "Check now" endpoint assembled the check transition through its own code path: the result record and the monitor state were written in two separate statements, a slow check could overlap the scheduled one for the same monitor, and the maintenance gate and notifications were skipped entirely -- a manual success after failures cleared the failing state and permanently swallowed the recovery notification. Both callers now execute one shared transition: record and state commit atomically, a concurrent check for the same monitor is refused with `409 Conflict`, and failure/recovery notifications fire under the same maintenance-window rules regardless of who triggered the check + +- **Temporary capability grants now bound running work** -- A temporary `terminal` grant's expiry blocked new sessions but left live PTY sessions running indefinitely, and the security-events pipeline was started once at boot from the permanent capability set, so a temporary `security_events` grant could never start it (nor could a persisted grant enable it across a restart) and expiry never stopped it. Capability state is now owned by one process-wide authority on the agent: expiry or revocation closes live terminal sessions immediately, and the security pipeline starts and stops as the capability becomes effective or lapses + ## [1.0.0-alpha.11] - 2026-07-03 ### Security diff --git a/apps/docs/content/docs/en/capabilities.mdx b/apps/docs/content/docs/en/capabilities.mdx index 8a76214a7..4af5d9201 100644 --- a/apps/docs/content/docs/en/capabilities.mdx +++ b/apps/docs/content/docs/en/capabilities.mdx @@ -138,6 +138,7 @@ Grants are persisted to `/capability_grants.json` (default `/var/lib/ - **Survives restarts** — if the agent restarts mid-window, the grant is reloaded and remains active for the rest of its original window. - **Expires at the original deadline** — restarting does not extend a grant. A grant for `30m` expires 30 minutes after it was issued regardless of how many times the agent restarts in between. - **Fails safe to OFF** — if the grants file is missing, empty, corrupt, or written by an unknown schema version, the agent treats it as "no grants" and the capability stays off. A temporary grant can never *strengthen* the permanent set; it can only flip an otherwise-off bit on for its window. +- **Expiry bounds running work, not just new requests** — when a `terminal` grant expires or is revoked, live terminal sessions opened under it are closed immediately (the browser sees the session end), and a `security_events` grant starts the agent's security pipeline the moment it becomes effective and stops it when the window closes — including grants issued while the agent was already running. ### What the server sees diff --git a/apps/docs/content/docs/zh/capabilities.mdx b/apps/docs/content/docs/zh/capabilities.mdx index 6be18ba85..a3d8aaddd 100644 --- a/apps/docs/content/docs/zh/capabilities.mdx +++ b/apps/docs/content/docs/zh/capabilities.mdx @@ -138,6 +138,7 @@ serverbee-agent grants - **跨重启存活** —— 若 Agent 在时间窗中重启,授予会被重新加载,并在原始时间窗内保持生效。 - **在原始截止时刻过期** —— 重启不会延长授予。`30m` 的授予在发起后 30 分钟过期,无论期间 Agent 重启多少次。 - **失败时安全归零(OFF)** —— 若 grants 文件缺失、为空、损坏,或由未知 schema 版本写入,Agent 会将其视作「无授予」,对应能力保持关闭。临时授予永远不会*增强*永久能力集,它只能在自己的时间窗内把一个原本关闭的位翻为开启。 +- **过期约束的是运行中的工作,而不只是新请求** —— `terminal` 授予过期或被撤销时,在该授予下打开的终端会话会被立即关闭(浏览器端会看到会话结束);`security_events` 授予生效的那一刻就会启动 Agent 的安全事件管道,时间窗结束时将其停止 —— 包括在 Agent 已运行期间发起的授予。 ### Server 看到什么 From c4d9efee5e4727c60714c23c95bbc491b0bd5fb9 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 12:36:04 +0800 Subject: [PATCH 05/44] refactor(server): own scheduled task lifecycle in one module Scheduled-task validation, persistence, in-memory jobs, execution, and startup restoration were split across service, task, and router modules. The service and task halves called each other, routers compensated database and scheduler mutations by hand, and cleanup paths could bypass the scheduler entirely. Deepen service::task_scheduler into the single lifecycle owner. Its interface now covers typed create/update inputs, lifecycle mutations, manual execution, startup restoration, and the two-phase server-cleanup plan; the router keeps only HTTP mapping and definition-level audit logs. Remove the old task module and the redundant cron dependency. Lifecycle fixes riding on the seam: - validate and calculate next runs with the scheduler parser, including correct weekday semantics - serialize lifecycle changes and fence queued callbacks by their registered job UUID - retain active-run ownership through cancellation, wait for completion before deletion, and reject stale cleanup - keep database rows, job registrations, orphan cleanup, and late task results synchronized across failures and commits --- Cargo.lock | 12 - crates/server/Cargo.toml | 1 - crates/server/src/main.rs | 3 +- crates/server/src/openapi.rs | 4 +- crates/server/src/router/api/server.rs | 22 +- crates/server/src/router/api/task.rs | 395 +--- crates/server/src/router/ws/browser.rs | 1 - crates/server/src/service/task_scheduler.rs | 1862 ++++++++++++++++++- crates/server/src/task/mod.rs | 1 - crates/server/src/task/task_scheduler.rs | 808 -------- 10 files changed, 1831 insertions(+), 1278 deletions(-) delete mode 100644 crates/server/src/task/task_scheduler.rs diff --git a/Cargo.lock b/Cargo.lock index f7c31ca29..a17c9ed73 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -794,17 +794,6 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "790eea4361631c5e7d22598ecd5723ff611904e3344ce8720784c93e3d83d40b" -[[package]] -name = "cron" -version = "0.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eee8b2b4516038bc0f1d3c9934bcb4a13dd316e04abbc63c96757a6d75978532" -dependencies = [ - "chrono", - "nom", - "once_cell", -] - [[package]] name = "croner" version = "2.2.0" @@ -4039,7 +4028,6 @@ dependencies = [ "base64 0.22.1", "chrono", "chrono-tz", - "cron", "dashmap", "dns-lookup", "figment", diff --git a/crates/server/Cargo.toml b/crates/server/Cargo.toml index 14166f744..28d974b2d 100644 --- a/crates/server/Cargo.toml +++ b/crates/server/Cargo.toml @@ -32,7 +32,6 @@ hickory-resolver = "0.25" utoipa = { version = "5", features = ["axum_extras", "chrono", "uuid"] } utoipa-swagger-ui = { version = "9", features = ["axum"] } tokio-cron-scheduler = "0.13" -cron = "0.13" async-trait = "0.1" futures-util = "0.3" once_cell = "1" diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs index 639d67b82..cc6aa0915 100644 --- a/crates/server/src/main.rs +++ b/crates/server/src/main.rs @@ -11,6 +11,7 @@ use serverbee_server::dev_demo; use serverbee_server::migration::Migrator; use serverbee_server::router::create_router; use serverbee_server::service::auth::AuthService; +use serverbee_server::service::task_scheduler; use serverbee_server::state::AppState; use serverbee_server::task; @@ -117,7 +118,7 @@ async fn main() -> anyhow::Result<()> { let s = state.clone(); tokio::spawn(async move { task::alert_evaluator::run(s).await }); let s = state.clone(); - tokio::spawn(async move { task::task_scheduler::run(s).await }); + tokio::spawn(async move { task_scheduler::restore_and_start(s).await }); let s = state.clone(); tokio::spawn(async move { task::service_monitor_checker::run(s).await }); let s = state.clone(); diff --git a/crates/server/src/openapi.rs b/crates/server/src/openapi.rs index c8d8eab67..cd2ca0be5 100644 --- a/crates/server/src/openapi.rs +++ b/crates/server/src/openapi.rs @@ -319,8 +319,8 @@ impl Modify for SecurityAddon { crate::service::alert::AlertEventResponse, crate::router::api::alert::AlertEventDetailResponse, // tasks - crate::router::api::task::CreateTaskRequest, - crate::router::api::task::UpdateTaskRequest, + crate::service::task_scheduler::CreateTaskRequest, + crate::service::task_scheduler::UpdateTaskRequest, crate::router::api::task::TaskResponse, // dashboards crate::entity::dashboard::Model, diff --git a/crates/server/src/router/api/server.rs b/crates/server/src/router/api/server.rs index 1c38e9eee..159817e46 100644 --- a/crates/server/src/router/api/server.rs +++ b/crates/server/src/router/api/server.rs @@ -29,6 +29,7 @@ use crate::service::network_probe::NetworkProbeService; use crate::service::record::{QueryHistoryResult, RecordService}; use crate::service::server::{ServerService, UpdateServerInput}; use crate::service::server_tag as server_tag_service; +use crate::service::task_scheduler; use crate::service::upgrade_tracker::{StartUpgradeJobError, UpgradeLookup}; use crate::state::AppState; use serverbee_common::protocol::ServerMessage; @@ -1235,6 +1236,7 @@ async fn cleanup_orphaned_servers( ) -> Result>, AppError> { use crate::entity::*; + let mut task_cleanup = task_scheduler::begin_server_cleanup(&state).await; let txn = state.db.begin().await?; let candidates = server::Entity::find() @@ -1256,7 +1258,11 @@ async fn cleanup_orphaned_servers( .exec(&txn) .await?; - // Tables with server_ids_json — per-table rules + task_cleanup + .remove_server_references(&txn, &orphan_ids) + .await?; + + // Remaining tables with server_ids_json — per-table rules cleanup_json_array_tables(&txn, &orphan_ids).await?; let deleted = server::Entity::delete_many() @@ -1265,6 +1271,7 @@ async fn cleanup_orphaned_servers( .await?; txn.commit().await?; + task_cleanup.apply_after_commit(&state).await; tracing::info!("Cleaned up {} orphaned servers", deleted.rows_affected); ok(CleanupResponse { @@ -1291,19 +1298,6 @@ async fn cleanup_json_array_tables( } } - // tasks: delete if empty - for t in task::Entity::find().all(txn).await? { - if let Some(new_json) = remove_ids_from_json(&t.server_ids_json, orphan_ids) { - if new_json == "[]" { - task::Entity::delete_by_id(&t.id).exec(txn).await?; - } else { - let mut active: task::ActiveModel = t.into(); - active.server_ids_json = Set(new_json); - active.update(txn).await?; - } - } - } - // alert_rules: delete if empty (+ related alert_states) for rule in alert_rule::Entity::find().all(txn).await? { if let Some(ref json) = rule.server_ids_json diff --git a/crates/server/src/router/api/task.rs b/crates/server/src/router/api/task.rs index 68b77dab0..6fb991e9c 100644 --- a/crates/server/src/router/api/task.rs +++ b/crates/server/src/router/api/task.rs @@ -1,4 +1,3 @@ -use std::str::FromStr; use std::sync::Arc; use axum::extract::{ConnectInfo, Extension, Path, Query, State}; @@ -8,17 +7,17 @@ use axum::{Json, Router}; use chrono::Utc; use sea_orm::*; use serde::{Deserialize, Serialize}; -use uuid::Uuid; -use crate::entity::{server, task, task_result}; +use crate::entity::{task, task_result}; use crate::error::{ApiResponse, AppError, ok}; use crate::middleware::auth::CurrentUser; use crate::router::utils::extract_client_ip; use crate::service::audit::AuditService; use crate::service::high_risk_audit::ExecAuditContext; +use crate::service::task_scheduler::{ + self as task_lifecycle, CreateTaskRequest, UpdateTaskRequest, +}; use crate::state::AppState; -use serverbee_common::constants::CAP_EXEC; -use serverbee_common::protocol::ServerMessage; pub fn router() -> Router> { Router::new() @@ -31,41 +30,6 @@ pub fn router() -> Router> { .route("/tasks/{id}/run", post(run_task)) } -// --- Request / Response types --- - -#[derive(Debug, Deserialize, utoipa::ToSchema)] -pub struct CreateTaskRequest { - pub command: String, - pub server_ids: Vec, - #[serde(default)] - pub timeout: Option, - /// "oneshot" (default) or "scheduled" - #[serde(default = "default_oneshot")] - pub task_type: String, - pub name: Option, - pub cron_expression: Option, - #[serde(default)] - pub retry_count: Option, - #[serde(default)] - pub retry_interval: Option, -} - -fn default_oneshot() -> String { - "oneshot".to_string() -} - -#[derive(Debug, Deserialize, utoipa::ToSchema)] -pub struct UpdateTaskRequest { - pub name: Option, - pub command: Option, - pub server_ids: Option>, - pub cron_expression: Option, - pub enabled: Option, - pub timeout: Option, - pub retry_count: Option, - pub retry_interval: Option, -} - #[derive(Debug, Deserialize, utoipa::IntoParams)] pub struct ListTasksQuery { #[serde(rename = "type")] @@ -156,119 +120,24 @@ pub async fn create_task( headers: HeaderMap, Json(input): Json, ) -> Result>, AppError> { - if input.server_ids.is_empty() { - return Err(AppError::Validation( - "server_ids cannot be empty".to_string(), - )); - } - if input.command.trim().is_empty() { - return Err(AppError::Validation("command cannot be empty".to_string())); - } - - let is_scheduled = input.task_type == "scheduled"; - - if is_scheduled { - let cron = input.cron_expression.as_deref().ok_or_else(|| { - AppError::Validation("cron_expression is required for scheduled tasks".into()) - })?; - // Validate cron expression - cron::Schedule::from_str(cron) - .map_err(|e| AppError::Validation(format!("Invalid cron expression: {e}")))?; - } - - // Validate numeric bounds - if matches!(input.timeout, Some(0)) { - return Err(AppError::Validation("timeout must be > 0".into())); - } - if let Some(rc) = input.retry_count - && !(0..=10).contains(&rc) - { - return Err(AppError::Validation( - "retry_count must be between 0 and 10".into(), - )); - } - if matches!(input.retry_interval, Some(ri) if ri < 1) { - return Err(AppError::Validation("retry_interval must be >= 1".into())); - } - - let task_id = Uuid::new_v4().to_string(); - let now = Utc::now(); let ip = extract_client_ip( &ConnectInfo(addr), &headers, &state.config.server.trusted_proxies, ) .to_string(); - let server_ids_json = serde_json::to_string(&input.server_ids) - .map_err(|e| AppError::Internal(format!("Serialization error: {e}")))?; - - // Compute next_run_at using the configured scheduler timezone - let tz: chrono_tz::Tz = state - .task_scheduler - .timezone() - .parse() - .unwrap_or(chrono_tz::UTC); - let next_run = input.cron_expression.as_deref().and_then(|c| { - cron::Schedule::from_str(c) - .ok() - .and_then(|s| s.upcoming(tz).next().map(|dt| dt.with_timezone(&Utc))) - }); - - let new_task = task::ActiveModel { - id: Set(task_id.clone()), - command: Set(input.command.clone()), - server_ids_json: Set(server_ids_json), - created_by: Set(current_user.user_id.clone()), - task_type: Set(input.task_type.clone()), - name: Set(input.name.clone()), - cron_expression: Set(input.cron_expression.clone()), - enabled: Set(true), - timeout: Set(input.timeout.map(|t| t as i32)), - retry_count: Set(input.retry_count.unwrap_or(0)), - retry_interval: Set(input.retry_interval.unwrap_or(60)), - last_run_at: NotSet, - next_run_at: Set(next_run), - created_at: Set(now), - }; - let task_model = new_task.insert(&state.db).await?; - - if is_scheduled { - // Register cron job in scheduler; rollback DB row on failure - if let Err(e) = state - .task_scheduler - .add_job(&task_model, state.clone()) - .await - { - let _ = task::Entity::delete_by_id(&task_id).exec(&state.db).await; - return Err(e); - } - tracing::info!("Scheduled task {} registered", task_id); - } else { - // One-shot: dispatch immediately - dispatch_oneshot( - &state, - &task_id, - &input.command, - &input.server_ids, - input.timeout, - ¤t_user.user_id, - &ip, - ) - .await?; - } - - // Audit task creation (best-effort). One-shot dispatch audits exec - // separately; this records the task definition itself for both kinds. + let audit_detail = format!( + "type={} name={} command={}", + input.task_type, + input.name.as_deref().unwrap_or("?"), + input.command + ); + let task_model = task_lifecycle::create_task(&state, input, ¤t_user.user_id, &ip).await?; let _ = AuditService::log( &state.db, ¤t_user.user_id, "task_created", - Some(&format!( - "task_id={task_id} type={} name={} command={}", - input.task_type, - input.name.as_deref().unwrap_or("?"), - input.command - )), + Some(&format!("task_id={} {audit_detail}", task_model.id)), &ip, ) .await; @@ -318,97 +187,7 @@ pub async fn update_task( Path(id): Path, Json(input): Json, ) -> Result>, AppError> { - let existing = task::Entity::find_by_id(&id) - .one(&state.db) - .await? - .ok_or_else(|| AppError::NotFound(format!("Task {id} not found")))?; - - let existing_backup = existing.clone(); - let mut model: task::ActiveModel = existing.into(); - - if let Some(name) = input.name { - model.name = Set(Some(name)); - } - if let Some(command) = input.command { - model.command = Set(command); - } - if let Some(server_ids) = &input.server_ids { - let json = serde_json::to_string(server_ids) - .map_err(|e| AppError::Internal(format!("Serialization error: {e}")))?; - model.server_ids_json = Set(json); - } - if let Some(cron) = &input.cron_expression { - cron::Schedule::from_str(cron) - .map_err(|e| AppError::Validation(format!("Invalid cron expression: {e}")))?; - // Recompute next_run_at using configured timezone - let tz: chrono_tz::Tz = state - .task_scheduler - .timezone() - .parse() - .unwrap_or(chrono_tz::UTC); - let next = cron::Schedule::from_str(cron) - .ok() - .and_then(|s| s.upcoming(tz).next().map(|dt| dt.with_timezone(&Utc))); - model.cron_expression = Set(Some(cron.clone())); - model.next_run_at = Set(next); - } - if let Some(enabled) = input.enabled { - model.enabled = Set(enabled); - // When resuming a paused task, recompute next_run_at so the UI shows a fresh time - if enabled { - let cron_expr = input - .cron_expression - .as_deref() - .or(existing_backup.cron_expression.as_deref()); - if let Some(cron) = cron_expr { - let tz: chrono_tz::Tz = state - .task_scheduler - .timezone() - .parse() - .unwrap_or(chrono_tz::UTC); - let next = cron::Schedule::from_str(cron) - .ok() - .and_then(|s| s.upcoming(tz).next().map(|dt| dt.with_timezone(&Utc))); - model.next_run_at = Set(next); - } - } - } - if let Some(timeout) = input.timeout { - if timeout < 1 { - return Err(AppError::Validation("timeout must be >= 1".into())); - } - model.timeout = Set(Some(timeout)); - } - if let Some(retry_count) = input.retry_count { - if !(0..=10).contains(&retry_count) { - return Err(AppError::Validation( - "retry_count must be between 0 and 10".into(), - )); - } - model.retry_count = Set(retry_count); - } - if let Some(retry_interval) = input.retry_interval { - if retry_interval < 1 { - return Err(AppError::Validation("retry_interval must be >= 1".into())); - } - model.retry_interval = Set(retry_interval); - } - - let updated = model.update(&state.db).await?; - - // Sync scheduler; on failure, restore the original row - if updated.task_type == "scheduled" - && let Err(e) = state - .task_scheduler - .update_job(&updated, state.clone()) - .await - { - let rollback: task::ActiveModel = existing_backup.into(); - let _ = rollback.update(&state.db).await; - return Err(e); - } - - // Audit the update (best-effort) once the row and scheduler are in sync. + let updated = task_lifecycle::update_task(&state, &id, input).await?; let ip = extract_client_ip( &ConnectInfo(addr), &headers, @@ -444,17 +223,7 @@ pub async fn delete_task( headers: HeaderMap, Path(id): Path, ) -> Result>, AppError> { - // Cancel active run and remove from scheduler first (idempotent, safe to call even if not registered) - state.task_scheduler.remove_job(&id).await?; - // Delete DB rows — if this fails, the scheduler job is already gone (acceptable: - // next server restart will simply not re-register the task since it was removed from scheduler) - task_result::Entity::delete_many() - .filter(task_result::Column::TaskId.eq(&id)) - .exec(&state.db) - .await?; - task::Entity::delete_by_id(&id).exec(&state.db).await?; - - // Audit the deletion (best-effort) after the rows are gone. + task_lifecycle::delete_task(&state, &id).await?; let ip = extract_client_ip( &ConnectInfo(addr), &headers, @@ -491,45 +260,22 @@ pub async fn run_task( headers: HeaderMap, Path(id): Path, ) -> Result>, AppError> { - // Validate task exists and is scheduled type - let task_model = task::Entity::find_by_id(&id) - .one(&state.db) - .await? - .ok_or_else(|| AppError::NotFound(format!("Task {id} not found")))?; - - if task_model.task_type != "scheduled" { - return Err(AppError::BadRequest( - "Only scheduled tasks can be manually triggered".into(), - )); - } - let ip = extract_client_ip( &ConnectInfo(addr), &headers, &state.config.server.trusted_proxies, ) .to_string(); - let started = crate::task::task_scheduler::execute_scheduled_task( + let updated = task_lifecycle::run_now( &state, &id, - true, Some(ExecAuditContext { user_id: current_user.user_id.clone(), ip, }), ) .await; - if !started { - return Err(AppError::Conflict( - "Task is currently running, try again later".into(), - )); - } - // Re-fetch to get updated last_run_at - let updated = task::Entity::find_by_id(&id) - .one(&state.db) - .await? - .ok_or_else(|| AppError::NotFound("Task not found".into()))?; - ok(updated.into()) + ok(updated?.into()) } #[utoipa::path( @@ -555,107 +301,6 @@ pub async fn get_task_results( ok(results) } -// --- Helper --- - -async fn dispatch_oneshot( - state: &Arc, - task_id: &str, - command: &str, - server_ids: &[String], - timeout: Option, - user_id: &str, - ip: &str, -) -> Result<(), AppError> { - let servers = server::Entity::find() - .filter(server::Column::Id.is_in(server_ids.iter().cloned())) - .all(&state.db) - .await?; - - let mut capable = Vec::new(); - let mut disabled = Vec::new(); - for sid in server_ids { - let configured_caps = servers - .iter() - .find(|s| s.id == *sid) - .map(|s| s.capabilities as u32) - .unwrap_or(0); - - if let Some(reason) = - state - .agent_manager - .capability_denied_reason(sid, configured_caps, CAP_EXEC) - { - disabled.push((sid.as_str(), exec_capability_denied_output(reason), reason)); - } else { - capable.push(sid); - } - } - - let now = Utc::now(); - for (sid, output, deny_reason) in &disabled { - let result = task_result::ActiveModel { - id: NotSet, - task_id: Set(task_id.to_string()), - server_id: Set((*sid).to_string()), - output: Set((*output).to_string()), - exit_code: Set(-2), - run_id: Set(None), - attempt: Set(1), - started_at: Set(None), - finished_at: Set(now), - }; - result.insert(&state.db).await?; - let detail = serde_json::json!({ - "server_id": sid, - "task_id": task_id, - "command": command, - "deny_reason": deny_reason, - }) - .to_string(); - let _ = AuditService::log(&state.db, user_id, "exec_denied", Some(&detail), ip).await; - } - - let mut dispatched = 0; - for sid in &capable { - if let Some(tx) = state.agent_manager.get_sender(sid) { - let msg = ServerMessage::Exec { - task_id: task_id.to_string(), - command: command.to_string(), - timeout, - }; - if tx.send(msg).await.is_ok() { - dispatched += 1; - let detail = serde_json::json!({ - "server_id": sid, - "task_id": task_id, - "command": command, - "timeout": timeout, - }) - .to_string(); - let _ = - AuditService::log(&state.db, user_id, "exec_started", Some(&detail), ip).await; - } - } - } - - tracing::info!( - "Task {} dispatched to {}/{} agents", - task_id, - dispatched, - server_ids.len() - ); - Ok(()) -} - -pub(crate) fn exec_capability_denied_output(reason: &str) -> &'static str { - match reason { - "agent_capability_disabled" => { - "Capability denied: exec is disabled in the agent's config (capabilities are agent-owned)" - } - _ => "Capability denied: exec disabled", - } -} - #[cfg(test)] mod audit_tests { use super::*; @@ -701,7 +346,9 @@ mod audit_tests { async fn delete_task_writes_audit_log() { let (db, _tmp) = setup_test_db().await; insert_oneshot_task(&db, "task-del").await; - let state = AppState::new(db.clone(), AppConfig::default()).await.unwrap(); + let state = AppState::new(db.clone(), AppConfig::default()) + .await + .unwrap(); let res = delete_task( State(state.clone()), @@ -726,7 +373,9 @@ mod audit_tests { async fn update_task_writes_audit_log() { let (db, _tmp) = setup_test_db().await; insert_oneshot_task(&db, "task-upd").await; - let state = AppState::new(db.clone(), AppConfig::default()).await.unwrap(); + let state = AppState::new(db.clone(), AppConfig::default()) + .await + .unwrap(); let res = update_task( State(state.clone()), diff --git a/crates/server/src/router/ws/browser.rs b/crates/server/src/router/ws/browser.rs index f2c905c62..4e7b7bf46 100644 --- a/crates/server/src/router/ws/browser.rs +++ b/crates/server/src/router/ws/browser.rs @@ -354,4 +354,3 @@ async fn send_browser_message( let text = serde_json::to_string(msg).map_err(axum::Error::new)?; sink.send(Message::Text(text.into())).await } - diff --git a/crates/server/src/service/task_scheduler.rs b/crates/server/src/service/task_scheduler.rs index de380518f..c2aea0deb 100644 --- a/crates/server/src/service/task_scheduler.rs +++ b/crates/server/src/service/task_scheduler.rs @@ -1,21 +1,127 @@ +use std::collections::HashMap; +use std::fmt; use std::sync::Arc; +use chrono::{DateTime, Utc}; use dashmap::DashMap; +use dashmap::mapref::entry::Entry; +use sea_orm::prelude::Expr; +use sea_orm::*; +use serde::Deserialize; +use tokio::sync::{Mutex, OwnedMutexGuard}; +use tokio::task::JoinSet; use tokio_cron_scheduler::{Job, JobScheduler}; use tokio_util::sync::CancellationToken; +use uuid::Uuid; +use crate::entity::{server, task, task_result}; use crate::error::AppError; +use crate::service::agent_manager::AgentRequestError; +use crate::service::audit::AuditService; +use crate::service::high_risk_audit::ExecAuditContext; +use crate::state::AppState; +use serverbee_common::constants::CAP_EXEC; +use serverbee_common::protocol::{AgentMessage, ServerMessage}; + +#[derive(Debug, Deserialize, utoipa::ToSchema)] +pub struct CreateTaskRequest { + pub command: String, + pub server_ids: Vec, + #[serde(default)] + pub timeout: Option, + /// "oneshot" (default) or "scheduled" + #[serde(default = "default_oneshot")] + pub task_type: TaskType, + pub name: Option, + pub cron_expression: Option, + #[serde(default)] + pub retry_count: Option, + #[serde(default)] + pub retry_interval: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, utoipa::ToSchema)] +#[serde(rename_all = "lowercase")] +pub enum TaskType { + Oneshot, + Scheduled, +} + +impl TaskType { + fn as_str(self) -> &'static str { + match self { + Self::Oneshot => "oneshot", + Self::Scheduled => "scheduled", + } + } +} + +impl fmt::Display for TaskType { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(self.as_str()) + } +} + +fn default_oneshot() -> TaskType { + TaskType::Oneshot +} + +#[derive(Debug, Deserialize, utoipa::ToSchema)] +pub struct UpdateTaskRequest { + pub name: Option, + pub command: Option, + pub server_ids: Option>, + pub cron_expression: Option, + pub enabled: Option, + pub timeout: Option, + pub retry_count: Option, + pub retry_interval: Option, +} + +struct ActiveRun { + run_id: String, + cancellation: CancellationToken, + completed: CancellationToken, +} + +type ActiveRuns = DashMap; + +fn clear_active_run_if_current(active_runs: &ActiveRuns, task_id: &str, run_id: &str) -> bool { + active_runs + .remove_if(task_id, |_, active_run| active_run.run_id == run_id) + .is_some() +} + +struct ActiveRunGuard { + active_runs: Arc, + task_id: String, + state: Arc, + run_id: String, + completed: CancellationToken, +} + +impl Drop for ActiveRunGuard { + fn drop(&mut self) { + clear_active_run_if_current(&self.active_runs, &self.task_id, &self.run_id); + self.state.exec_audit_contexts.remove(&self.run_id); + self.completed.cancel(); + } +} pub struct TaskScheduler { scheduler: JobScheduler, job_map: DashMap, - /// task_id -> (run_id, cancellation token). Arc so the cleanup guard shares the real map. - pub(crate) active_runs: Arc>, - timezone: String, + /// Active task executions. Arc so cleanup guards share the real map. + active_runs: Arc, + lifecycle_lock: Arc>, + timezone: chrono_tz::Tz, } impl TaskScheduler { pub async fn new(timezone: &str) -> Result { + let timezone = timezone + .parse() + .map_err(|_| AppError::Internal(format!("Invalid timezone: {timezone}")))?; let scheduler = JobScheduler::new() .await .map_err(|e| AppError::Internal(format!("Failed to create scheduler: {e}")))?; @@ -23,25 +129,46 @@ impl TaskScheduler { scheduler, job_map: DashMap::new(), active_runs: Arc::new(DashMap::new()), - timezone: timezone.to_string(), + lifecycle_lock: Arc::new(Mutex::new(())), + timezone, }) } - pub fn is_running(&self, task_id: &str) -> bool { + #[cfg(test)] + fn is_running(&self, task_id: &str) -> bool { self.active_runs.contains_key(task_id) } - pub fn cancel_active_run(&self, task_id: &str) { - if let Some((_, (_, token))) = self.active_runs.remove(task_id) { - token.cancel(); + fn cancel_active_run(&self, task_id: &str) -> Option { + if let Some(active_run) = self.active_runs.get(task_id) { + active_run.cancellation.cancel(); + return Some(active_run.completed.clone()); + } + None + } + + async fn cancel_and_wait_active_run(&self, task_id: &str) { + if let Some(completed) = self.cancel_active_run(task_id) { + completed.cancelled().await; } } - pub fn timezone(&self) -> &str { - &self.timezone + fn tz(&self) -> chrono_tz::Tz { + self.timezone + } + + fn next_run_at(&self, cron_expr: &str) -> Option> { + let mut job = Job::new_tz(cron_expr, self.tz(), |_uuid, _lock| {}).ok()?; + job.job_data().ok()?.next_tick_utc() } - pub async fn start(&self) -> Result<(), AppError> { + fn is_current_job(&self, task_id: &str, job_id: uuid::Uuid) -> bool { + self.job_map + .get(task_id) + .is_some_and(|current_job_id| *current_job_id == job_id) + } + + async fn start(&self) -> Result<(), AppError> { self.scheduler .start() .await @@ -49,28 +176,45 @@ impl TaskScheduler { Ok(()) } - pub async fn add_job( + async fn add_job( &self, - task_model: &crate::entity::task::Model, - state: std::sync::Arc, + task_model: &task::Model, + state: Arc, ) -> Result<(), AppError> { + if self.job_map.contains_key(&task_model.id) { + return Err(AppError::Internal(format!( + "Scheduled task {} is already registered", + task_model.id + ))); + } let cron = task_model .cron_expression .as_deref() .ok_or_else(|| AppError::BadRequest("Missing cron_expression".into()))?; let task_id = task_model.id.clone(); - let tz: chrono_tz::Tz = self - .timezone - .parse() - .map_err(|_| AppError::Internal(format!("Invalid timezone: {}", self.timezone)))?; - - let job = Job::new_async_tz(cron, tz, move |_uuid, _lock| { + let job = Job::new_async_tz(cron, self.tz(), move |job_id, _lock| { let state = state.clone(); let task_id = task_id.clone(); Box::pin(async move { - crate::task::task_scheduler::execute_scheduled_task(&state, &task_id, false, None) - .await; + let _lifecycle_guard = state.task_scheduler.lifecycle_lock.lock().await; + if !state.task_scheduler.is_current_job(&task_id, job_id) { + tracing::debug!( + task_id, + %job_id, + "Ignoring stale scheduled task callback" + ); + return; + } + match execute_scheduled_task(&state, &task_id, false, None).await { + Ok(true) => {} + Ok(false) => { + tracing::warn!("Task {task_id} still running, skipping cron trigger"); + } + Err(error) => { + tracing::error!("Failed to start scheduled task {task_id}: {error}"); + } + } }) }) .map_err(|e| AppError::BadRequest(format!("Invalid cron expression: {e}")))?; @@ -84,21 +228,29 @@ impl TaskScheduler { Ok(()) } - pub async fn remove_job(&self, task_id: &str) -> Result<(), AppError> { - self.cancel_active_run(task_id); - if let Some((_, job_id)) = self.job_map.remove(task_id) { + async fn remove_job(&self, task_id: &str) -> Result<(), AppError> { + let job_id = self + .job_map + .get(task_id) + .map(|entry| entry.value().to_owned()); + let remove_result = if let Some(job_id) = job_id { self.scheduler .remove(&job_id) .await - .map_err(|e| AppError::Internal(format!("Failed to remove job: {e}")))?; - } + .map_err(|error| AppError::Internal(format!("Failed to remove job: {error}"))) + } else { + Ok(()) + }; + self.cancel_and_wait_active_run(task_id).await; + remove_result?; + self.job_map.remove(task_id); Ok(()) } - pub async fn update_job( + async fn update_job( &self, - task_model: &crate::entity::task::Model, - state: std::sync::Arc, + task_model: &task::Model, + state: Arc, ) -> Result<(), AppError> { self.remove_job(&task_model.id).await?; if task_model.enabled { @@ -107,14 +259,936 @@ impl TaskScheduler { Ok(()) } - pub async fn disable_job(&self, task_id: &str) -> Result<(), AppError> { - self.remove_job(task_id).await + async fn restore_job( + &self, + task_model: &task::Model, + state: Arc, + ) -> Result<(), AppError> { + if task_model.task_type == "scheduled" + && task_model.enabled + && !self.job_map.contains_key(&task_model.id) + { + self.add_job(task_model, state).await?; + } + Ok(()) + } +} + +pub fn validate_cron(expr: &str) -> Result<(), AppError> { + Job::new_tz(expr, chrono_tz::UTC, |_uuid, _lock| {}) + .map(|_| ()) + .map_err(|error| AppError::Validation(format!("Invalid cron expression: {error}"))) +} + +pub async fn create_task( + state: &Arc, + input: CreateTaskRequest, + created_by: &str, + ip: &str, +) -> Result { + let CreateTaskRequest { + command, + server_ids, + timeout, + task_type, + name, + cron_expression, + retry_count, + retry_interval, + } = input; + + if server_ids.is_empty() { + return Err(AppError::Validation( + "server_ids cannot be empty".to_string(), + )); + } + if command.trim().is_empty() { + return Err(AppError::Validation("command cannot be empty".to_string())); + } + + let is_scheduled = task_type == TaskType::Scheduled; + if is_scheduled { + let cron = cron_expression.as_deref().ok_or_else(|| { + AppError::Validation("cron_expression is required for scheduled tasks".into()) + })?; + validate_cron(cron)?; + } + + let timeout_i32 = match timeout { + Some(0) => return Err(AppError::Validation("timeout must be > 0".into())), + Some(value) => Some( + i32::try_from(value) + .map_err(|_| AppError::Validation("timeout exceeds supported range".into()))?, + ), + None => None, + }; + if let Some(count) = retry_count + && !(0..=10).contains(&count) + { + return Err(AppError::Validation( + "retry_count must be between 0 and 10".into(), + )); + } + if matches!(retry_interval, Some(interval) if interval < 1) { + return Err(AppError::Validation("retry_interval must be >= 1".into())); + } + + let _lifecycle_guard = if is_scheduled { + Some(state.task_scheduler.lifecycle_lock.lock().await) + } else { + None + }; + let task_id = Uuid::new_v4().to_string(); + let server_ids_json = serde_json::to_string(&server_ids) + .map_err(|error| AppError::Internal(format!("Serialization error: {error}")))?; + let next_run_at = cron_expression + .as_deref() + .and_then(|expr| state.task_scheduler.next_run_at(expr)); + let task_model = task::ActiveModel { + id: Set(task_id.clone()), + command: Set(command.clone()), + server_ids_json: Set(server_ids_json), + created_by: Set(created_by.to_string()), + task_type: Set(task_type.as_str().to_string()), + name: Set(name), + cron_expression: Set(cron_expression), + enabled: Set(true), + timeout: Set(timeout_i32), + retry_count: Set(retry_count.unwrap_or(0)), + retry_interval: Set(retry_interval.unwrap_or(60)), + last_run_at: NotSet, + next_run_at: Set(next_run_at), + created_at: Set(Utc::now()), + } + .insert(&state.db) + .await?; + + if is_scheduled { + // The row is persisted before the scheduler's asynchronous mutation. + // Keep both sides in sync if registration fails after validation. + if let Err(error) = state + .task_scheduler + .add_job(&task_model, state.clone()) + .await + { + if let Err(rollback_error) = task::Entity::delete_by_id(&task_id).exec(&state.db).await + { + tracing::error!( + task_id, + "Failed to remove task row after scheduler registration error: {rollback_error}" + ); + } + return Err(error); + } + tracing::info!("Scheduled task {} registered", task_id); + } else { + dispatch_oneshot( + state, + &task_id, + &command, + &server_ids, + timeout, + created_by, + ip, + ) + .await?; + } + + Ok(task_model) +} + +pub async fn update_task( + state: &Arc, + id: &str, + input: UpdateTaskRequest, +) -> Result { + let _lifecycle_guard = state.task_scheduler.lifecycle_lock.lock().await; + let existing = task::Entity::find_by_id(id) + .one(&state.db) + .await? + .ok_or_else(|| AppError::NotFound(format!("Task {id} not found")))?; + let existing_backup = existing.clone(); + let mut model: task::ActiveModel = existing.into(); + + let UpdateTaskRequest { + name, + command, + server_ids, + cron_expression, + enabled, + timeout, + retry_count, + retry_interval, + } = input; + + if let Some(name) = name { + model.name = Set(Some(name)); + } + if let Some(command) = command { + model.command = Set(command); + } + if let Some(server_ids) = server_ids { + let json = serde_json::to_string(&server_ids) + .map_err(|error| AppError::Internal(format!("Serialization error: {error}")))?; + model.server_ids_json = Set(json); + } + if let Some(cron) = cron_expression.as_deref() { + validate_cron(cron)?; + model.cron_expression = Set(Some(cron.to_string())); + model.next_run_at = Set(state.task_scheduler.next_run_at(cron)); + } + if let Some(enabled) = enabled { + model.enabled = Set(enabled); + if enabled { + let cron = cron_expression + .as_deref() + .or(existing_backup.cron_expression.as_deref()); + if let Some(cron) = cron { + model.next_run_at = Set(state.task_scheduler.next_run_at(cron)); + } + } + } + if let Some(timeout) = timeout { + if timeout < 1 { + return Err(AppError::Validation("timeout must be >= 1".into())); + } + model.timeout = Set(Some(timeout)); + } + if let Some(retry_count) = retry_count { + if !(0..=10).contains(&retry_count) { + return Err(AppError::Validation( + "retry_count must be between 0 and 10".into(), + )); + } + model.retry_count = Set(retry_count); + } + if let Some(retry_interval) = retry_interval { + if retry_interval < 1 { + return Err(AppError::Validation("retry_interval must be >= 1".into())); + } + model.retry_interval = Set(retry_interval); + } + + let updated = model.update(&state.db).await?; + if updated.task_type == "scheduled" + && let Err(error) = state + .task_scheduler + .update_job(&updated, state.clone()) + .await + { + let rollback: task::ActiveModel = existing_backup.clone().into(); + match rollback.update(&state.db).await { + Ok(restored) => { + if let Err(restore_error) = state + .task_scheduler + .restore_job(&restored, state.clone()) + .await + { + tracing::error!( + task_id = id, + "Failed to restore scheduled job after update error: {restore_error}" + ); + } + } + Err(rollback_error) => { + tracing::error!( + task_id = id, + "Failed to restore task row after scheduler update error: {rollback_error}" + ); + } + } + return Err(error); + } + + Ok(updated) +} + +pub async fn delete_task(state: &Arc, id: &str) -> Result<(), AppError> { + let _lifecycle_guard = state.task_scheduler.lifecycle_lock.lock().await; + let existing = task::Entity::find_by_id(id).one(&state.db).await?; + state.task_scheduler.remove_job(id).await?; + + let transaction = match state.db.begin().await { + Ok(transaction) => transaction, + Err(error) => { + if let Some(task_model) = existing.as_ref() + && let Err(restore_error) = state + .task_scheduler + .restore_job(task_model, state.clone()) + .await + { + tracing::error!( + task_id = id, + "Failed to restore scheduled job after transaction error: {restore_error}" + ); + } + return Err(error.into()); + } + }; + let delete_result: Result<(), DbErr> = async { + task_result::Entity::delete_many() + .filter(task_result::Column::TaskId.eq(id)) + .exec(&transaction) + .await?; + task::Entity::delete_by_id(id).exec(&transaction).await?; + transaction.commit().await?; + Ok(()) + } + .await; + + if let Err(error) = delete_result { + if let Some(task_model) = existing.as_ref() + && let Err(restore_error) = state + .task_scheduler + .restore_job(task_model, state.clone()) + .await + { + tracing::error!( + task_id = id, + "Failed to restore scheduled job after delete error: {restore_error}" + ); + } + return Err(error.into()); + } + + Ok(()) +} + +#[must_use = "the scheduler cleanup must run after the surrounding transaction commits"] +pub(crate) struct TaskServerCleanup { + changed_task_ids: Vec, + deleted_task_ids: Vec, + orphan_server_ids: Vec, + _lifecycle_guard: OwnedMutexGuard<()>, +} + +impl TaskServerCleanup { + pub(crate) async fn remove_server_references( + &mut self, + transaction: &DatabaseTransaction, + orphan_ids: &[String], + ) -> Result<(), AppError> { + self.orphan_server_ids = orphan_ids.to_vec(); + for task_model in task::Entity::find().all(transaction).await? { + let mut server_ids: Vec = serde_json::from_str(&task_model.server_ids_json) + .map_err(|error| { + AppError::Internal(format!( + "Task {} has invalid server_ids_json: {error}", + task_model.id + )) + })?; + let previous_len = server_ids.len(); + server_ids.retain(|server_id| !orphan_ids.contains(server_id)); + if server_ids.len() == previous_len { + continue; + } + self.changed_task_ids.push(task_model.id.clone()); + + if server_ids.is_empty() { + task_result::Entity::delete_many() + .filter(task_result::Column::TaskId.eq(&task_model.id)) + .exec(transaction) + .await?; + task::Entity::delete_by_id(&task_model.id) + .exec(transaction) + .await?; + self.deleted_task_ids.push(task_model.id); + } else { + let mut active: task::ActiveModel = task_model.into(); + active.server_ids_json = + Set(serde_json::to_string(&server_ids).map_err(|error| { + AppError::Internal(format!("Failed to serialize task server IDs: {error}")) + })?); + active.update(transaction).await?; + } + } + + Ok(()) + } + + pub(crate) async fn apply_after_commit(self, state: &Arc) { + let deleted_task_ids: std::collections::HashSet<&str> = + self.deleted_task_ids.iter().map(String::as_str).collect(); + for task_id in &self.changed_task_ids { + if deleted_task_ids.contains(task_id.as_str()) { + if let Err(error) = state.task_scheduler.remove_job(task_id).await { + tracing::error!( + task_id, + "Failed to remove cleaned-up task from scheduler: {error}" + ); + } + } else { + state + .task_scheduler + .cancel_and_wait_active_run(task_id) + .await; + } + } + + if !self.deleted_task_ids.is_empty() + && let Err(error) = task_result::Entity::delete_many() + .filter(task_result::Column::TaskId.is_in(self.deleted_task_ids.clone())) + .exec(&state.db) + .await + { + tracing::error!("Failed to remove late results for deleted tasks: {error}"); + } + if !self.orphan_server_ids.is_empty() + && let Err(error) = task_result::Entity::delete_many() + .filter(task_result::Column::ServerId.is_in(self.orphan_server_ids.clone())) + .exec(&state.db) + .await + { + tracing::error!("Failed to remove late results for cleaned-up servers: {error}"); + } + } +} + +pub(crate) async fn begin_server_cleanup(state: &Arc) -> TaskServerCleanup { + let lifecycle_guard = state + .task_scheduler + .lifecycle_lock + .clone() + .lock_owned() + .await; + TaskServerCleanup { + changed_task_ids: Vec::new(), + deleted_task_ids: Vec::new(), + orphan_server_ids: Vec::new(), + _lifecycle_guard: lifecycle_guard, + } +} + +pub async fn run_now( + state: &Arc, + id: &str, + audit_context: Option, +) -> Result { + let _lifecycle_guard = state.task_scheduler.lifecycle_lock.lock().await; + let task_model = task::Entity::find_by_id(id) + .one(&state.db) + .await? + .ok_or_else(|| AppError::NotFound(format!("Task {id} not found")))?; + if task_model.task_type != "scheduled" { + return Err(AppError::BadRequest( + "Only scheduled tasks can be manually triggered".into(), + )); + } + + if !execute_scheduled_task(state, id, true, audit_context).await? { + return Err(AppError::Conflict( + "Task is currently running, try again later".into(), + )); + } + + task::Entity::find_by_id(id) + .one(&state.db) + .await? + .ok_or_else(|| AppError::NotFound("Task not found".into())) +} + +/// Restore enabled scheduled tasks from the database, then start the scheduler. +pub async fn restore_and_start(state: Arc) { + let _lifecycle_guard = state.task_scheduler.lifecycle_lock.lock().await; + let tasks = task::Entity::find() + .filter(task::Column::TaskType.eq("scheduled")) + .filter(task::Column::Enabled.eq(true)) + .all(&state.db) + .await; + + match tasks { + Ok(tasks) => { + for task_model in &tasks { + if let Err(error) = state + .task_scheduler + .add_job(task_model, state.clone()) + .await + { + tracing::error!( + "Failed to register scheduled task {}: {error}", + task_model.id + ); + continue; + } + let next_run_at = task_model + .cron_expression + .as_deref() + .and_then(|expr| state.task_scheduler.next_run_at(expr)); + if let Err(error) = task::Entity::update_many() + .filter(task::Column::Id.eq(&task_model.id)) + .col_expr(task::Column::NextRunAt, Expr::value(next_run_at)) + .exec(&state.db) + .await + { + tracing::error!( + task_id = task_model.id, + "Failed to refresh next run during scheduler restore: {error}" + ); + } + } + tracing::info!("Loaded {} scheduled tasks", tasks.len()); + } + Err(error) => { + tracing::error!("Failed to load scheduled tasks: {error}"); + } + } + + if let Err(error) = state.task_scheduler.start().await { + tracing::error!("Failed to start task scheduler: {error}"); + } +} + +/// Build correlation ID: {task_id}:{run_id}:{server_id}:{attempt} +fn build_correlation_id(task_id: &str, run_id: &str, server_id: &str, attempt: i32) -> String { + format!("{task_id}:{run_id}:{server_id}:{attempt}") +} + +/// Called by a cron trigger or the manual run interface. +/// Returns true if execution was started, false if an overlapping run won. +async fn execute_scheduled_task( + state: &Arc, + task_id: &str, + skip_retry: bool, + audit_context: Option, +) -> Result { + let scheduler = &state.task_scheduler; + let run_id = Uuid::new_v4().to_string(); + let token = CancellationToken::new(); + let completed = CancellationToken::new(); + match scheduler.active_runs.entry(task_id.to_string()) { + Entry::Occupied(_) => { + tracing::warn!("Task {task_id} still running, skipping trigger"); + return Ok(false); + } + Entry::Vacant(entry) => { + entry.insert(ActiveRun { + run_id: run_id.clone(), + cancellation: token.clone(), + completed: completed.clone(), + }); + } + } + + let active_run_guard = ActiveRunGuard { + active_runs: Arc::clone(&scheduler.active_runs), + task_id: task_id.to_string(), + state: state.clone(), + run_id: run_id.clone(), + completed, + }; + let task_model = task::Entity::find_by_id(task_id) + .one(&state.db) + .await? + .ok_or_else(|| AppError::NotFound(format!("Task {task_id} not found")))?; + + let server_ids: Vec = serde_json::from_str(&task_model.server_ids_json) + .map_err(|error| AppError::Internal(format!("Invalid task server_ids_json: {error}")))?; + let timeout_secs = task_model.timeout.unwrap_or(300).max(1) as u64; + let retry_count = if skip_retry { + 0 + } else { + task_model.retry_count.max(0) + }; + let retry_interval = task_model.retry_interval.max(1) as u64; + let command = task_model.command.clone(); + let audit_context_ref = audit_context.as_ref(); + if let Some(context) = audit_context.clone() { + state.exec_audit_contexts.insert(run_id.clone(), context); + } + + let next_run_at = task_model + .cron_expression + .as_deref() + .and_then(|expr| scheduler.next_run_at(expr)); + let mut update = task::Entity::update_many() + .filter(task::Column::Id.eq(task_id)) + .col_expr(task::Column::LastRunAt, Expr::value(Utc::now())); + if let Some(next_run_at) = next_run_at { + update = update.col_expr(task::Column::NextRunAt, Expr::value(next_run_at)); + } + update.exec(&state.db).await?; + + let target_servers = server::Entity::find() + .filter(server::Column::Id.is_in(server_ids.clone())) + .all(&state.db) + .await?; + let server_capabilities: HashMap = target_servers + .iter() + .map(|server| (server.id.clone(), server.capabilities)) + .collect(); + let mut join_set = JoinSet::new(); + + for server_id in &server_ids { + let configured_capabilities = server_capabilities.get(server_id).copied().unwrap_or(0); + if let Some(reason) = state.agent_manager.capability_denied_reason( + server_id, + configured_capabilities as u32, + CAP_EXEC, + ) { + write_synthetic_result( + &state.db, + task_id, + &run_id, + server_id, + -2, + exec_capability_denied_output(reason), + ) + .await?; + if let Some(context) = audit_context_ref { + let detail = serde_json::json!({ + "server_id": server_id, + "task_id": task_id, + "command": command, + "deny_reason": reason, + }) + .to_string(); + let _ = AuditService::log( + &state.db, + &context.user_id, + "exec_denied", + Some(&detail), + &context.ip, + ) + .await; + } + continue; + } + + if let Some(context) = audit_context_ref { + let detail = serde_json::json!({ + "server_id": server_id, + "task_id": task_id, + "command": command, + "timeout": Some(timeout_secs as u32), + }) + .to_string(); + let _ = AuditService::log( + &state.db, + &context.user_id, + "exec_started", + Some(&detail), + &context.ip, + ) + .await; + } + + let state = state.clone(); + let task_id = task_id.to_string(); + let run_id = run_id.clone(); + let server_id = server_id.clone(); + let command = command.clone(); + let token = token.clone(); + join_set.spawn(async move { + execute_for_server( + &state, + &task_id, + &run_id, + &server_id, + &command, + timeout_secs, + retry_count, + retry_interval, + token, + ) + .await + }); + } + + tokio::spawn(async move { + let _active_run_guard = active_run_guard; + while let Some(result) = join_set.join_next().await { + match result { + Ok(Ok(())) => {} + Ok(Err(error)) => { + tracing::error!("Failed to persist scheduled task result: {error}"); + } + Err(error) => { + tracing::error!("Scheduled task executor failed to join: {error}"); + } + } + } + }); + + Ok(true) +} + +#[allow(clippy::too_many_arguments)] +async fn execute_for_server( + state: &Arc, + task_id: &str, + run_id: &str, + server_id: &str, + command: &str, + timeout_secs: u64, + retry_count: i32, + retry_interval: u64, + token: CancellationToken, +) -> Result<(), DbErr> { + let max_attempts = retry_count + 1; + let timeout_duration = std::time::Duration::from_secs(timeout_secs + 10); + + for attempt in 1..=max_attempts { + if token.is_cancelled() { + break; + } + if attempt > 1 { + tokio::select! { + _ = tokio::time::sleep(std::time::Duration::from_secs(retry_interval)) => {} + _ = token.cancelled() => { break; } + } + } + + let correlation_id = build_correlation_id(task_id, run_id, server_id, attempt); + let started_at = Utc::now(); + let result = tokio::select! { + result = state.agent_manager.request_with_id( + server_id, + correlation_id, + timeout_duration, + |msg_id| ServerMessage::Exec { + task_id: msg_id, + command: command.to_string(), + timeout: Some(timeout_secs as u32), + }, + ) => result, + _ = token.cancelled() => { + break; + } + }; + + if token.is_cancelled() { + break; + } + + match result { + Ok(AgentMessage::TaskResult { result, .. }) => { + write_result( + &state.db, + task_id, + run_id, + server_id, + attempt, + started_at, + result.exit_code, + &result.output, + ) + .await?; + if result.exit_code == 0 { + break; + } + } + Err(AgentRequestError::Offline) => { + write_result( + &state.db, + task_id, + run_id, + server_id, + attempt, + started_at, + -3, + "Server offline", + ) + .await?; + } + Err(AgentRequestError::SendFailed) => { + write_result( + &state.db, + task_id, + run_id, + server_id, + attempt, + started_at, + -3, + "Dispatch failed", + ) + .await?; + } + _ => { + write_result( + &state.db, + task_id, + run_id, + server_id, + attempt, + started_at, + -4, + &format!("No response within {timeout_secs}s"), + ) + .await?; + } + } + + if attempt == max_attempts { + break; + } + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +async fn write_result( + db: &DatabaseConnection, + task_id: &str, + run_id: &str, + server_id: &str, + attempt: i32, + started_at: DateTime, + exit_code: i32, + output: &str, +) -> Result<(), DbErr> { + task_result::ActiveModel { + id: NotSet, + task_id: Set(task_id.to_string()), + server_id: Set(server_id.to_string()), + output: Set(output.to_string()), + exit_code: Set(exit_code), + finished_at: Set(Utc::now()), + run_id: Set(Some(run_id.to_string())), + attempt: Set(attempt), + started_at: Set(Some(started_at)), + } + .insert(db) + .await?; + Ok(()) +} + +async fn write_synthetic_result( + db: &DatabaseConnection, + task_id: &str, + run_id: &str, + server_id: &str, + exit_code: i32, + output: &str, +) -> Result<(), DbErr> { + write_result( + db, + task_id, + run_id, + server_id, + 1, + Utc::now(), + exit_code, + output, + ) + .await +} + +async fn dispatch_oneshot( + state: &Arc, + task_id: &str, + command: &str, + server_ids: &[String], + timeout: Option, + user_id: &str, + ip: &str, +) -> Result<(), AppError> { + let servers = server::Entity::find() + .filter(server::Column::Id.is_in(server_ids.iter().cloned())) + .all(&state.db) + .await?; + let mut capable = Vec::new(); + let mut disabled = Vec::new(); + + for server_id in server_ids { + let configured_capabilities = servers + .iter() + .find(|server| server.id == *server_id) + .map(|server| server.capabilities as u32) + .unwrap_or(0); + if let Some(reason) = state.agent_manager.capability_denied_reason( + server_id, + configured_capabilities, + CAP_EXEC, + ) { + disabled.push(( + server_id.as_str(), + exec_capability_denied_output(reason), + reason, + )); + } else { + capable.push(server_id); + } + } + + let now = Utc::now(); + for (server_id, output, deny_reason) in &disabled { + task_result::ActiveModel { + id: NotSet, + task_id: Set(task_id.to_string()), + server_id: Set((*server_id).to_string()), + output: Set((*output).to_string()), + exit_code: Set(-2), + run_id: Set(None), + attempt: Set(1), + started_at: Set(None), + finished_at: Set(now), + } + .insert(&state.db) + .await?; + let detail = serde_json::json!({ + "server_id": server_id, + "task_id": task_id, + "command": command, + "deny_reason": deny_reason, + }) + .to_string(); + let _ = AuditService::log(&state.db, user_id, "exec_denied", Some(&detail), ip).await; + } + + let mut dispatched = 0; + for server_id in &capable { + if let Some(sender) = state.agent_manager.get_sender(server_id) { + let message = ServerMessage::Exec { + task_id: task_id.to_string(), + command: command.to_string(), + timeout, + }; + if sender.send(message).await.is_ok() { + dispatched += 1; + let detail = serde_json::json!({ + "server_id": server_id, + "task_id": task_id, + "command": command, + "timeout": timeout, + }) + .to_string(); + let _ = + AuditService::log(&state.db, user_id, "exec_started", Some(&detail), ip).await; + } + } + } + + tracing::info!( + "Task {} dispatched to {}/{} agents", + task_id, + dispatched, + server_ids.len() + ); + Ok(()) +} + +fn exec_capability_denied_output(reason: &str) -> &'static str { + match reason { + "agent_capability_disabled" => { + "Capability denied: exec is disabled in the agent's config (capabilities are agent-owned)" + } + _ => "Capability denied: exec disabled", } } #[cfg(test)] mod tests { use super::*; + use chrono::{Datelike, TimeZone, Weekday}; + use serverbee_common::constants::CAP_DEFAULT; + + fn test_active_run(run_id: &str) -> ActiveRun { + ActiveRun { + run_id: run_id.to_string(), + cancellation: CancellationToken::new(), + completed: CancellationToken::new(), + } + } #[tokio::test] async fn test_new_scheduler() { @@ -125,10 +1199,9 @@ mod tests { #[tokio::test] async fn test_overlap_detection() { let scheduler = TaskScheduler::new("UTC").await.unwrap(); - let token = CancellationToken::new(); scheduler .active_runs - .insert("task-1".to_string(), ("run-1".to_string(), token)); + .insert("task-1".to_string(), test_active_run("run-1")); assert!(scheduler.is_running("task-1")); assert!(!scheduler.is_running("task-2")); } @@ -136,14 +1209,47 @@ mod tests { #[tokio::test] async fn test_cancel_active_run() { let scheduler = TaskScheduler::new("UTC").await.unwrap(); - let token = CancellationToken::new(); - let token_clone = token.clone(); + let active_run = test_active_run("run-1"); + let token = active_run.cancellation.clone(); scheduler .active_runs - .insert("task-1".to_string(), ("run-1".to_string(), token)); + .insert("task-1".to_string(), active_run); scheduler.cancel_active_run("task-1"); - assert!(token_clone.is_cancelled()); - assert!(!scheduler.is_running("task-1")); + assert!(token.is_cancelled()); + assert!( + scheduler.is_running("task-1"), + "the slot must remain occupied until the cancelled run drains" + ); + } + + #[tokio::test] + async fn test_stale_cleanup_cannot_remove_a_newer_run() { + let scheduler = TaskScheduler::new("UTC").await.unwrap(); + scheduler + .active_runs + .insert("task-1".to_string(), test_active_run("run-a")); + scheduler.cancel_active_run("task-1"); + + assert!(matches!( + scheduler.active_runs.entry("task-1".to_string()), + Entry::Occupied(_) + )); + assert!(clear_active_run_if_current( + &scheduler.active_runs, + "task-1", + "run-a" + )); + + scheduler + .active_runs + .insert("task-1".to_string(), test_active_run("run-b")); + assert!(!clear_active_run_if_current( + &scheduler.active_runs, + "task-1", + "run-a" + )); + let active = scheduler.active_runs.get("task-1").unwrap(); + assert_eq!(active.run_id, "run-b"); } // ---- Helpers ------------------------------------------------------- @@ -194,12 +1300,25 @@ mod tests { } } + fn create_scheduled_request() -> CreateTaskRequest { + CreateTaskRequest { + command: "echo hi".to_string(), + server_ids: vec!["server-1".to_string()], + timeout: Some(5), + task_type: TaskType::Scheduled, + name: Some("test task".to_string()), + cron_expression: Some(FAR_FUTURE_CRON.to_string()), + retry_count: Some(0), + retry_interval: Some(1), + } + } + // ---- accessors / basic state -------------------------------------- #[tokio::test] async fn test_timezone_accessor() { let scheduler = TaskScheduler::new("Asia/Shanghai").await.unwrap(); - assert_eq!(scheduler.timezone(), "Asia/Shanghai"); + assert_eq!(scheduler.tz(), chrono_tz::Asia::Shanghai); } #[tokio::test] @@ -233,12 +1352,11 @@ mod tests { } #[tokio::test] - async fn test_add_job_invalid_timezone() { - let (state, _db, _dir) = build_test_state().await; - // Construction does not validate the tz; add_job does. - let scheduler = TaskScheduler::new("Not/AZone").await.unwrap(); - let task = make_task("t-bad-tz", Some(FAR_FUTURE_CRON), true); - let err = scheduler.add_job(&task, state).await.unwrap_err(); + async fn test_new_scheduler_rejects_invalid_timezone() { + let err = match TaskScheduler::new("Not/AZone").await { + Ok(_) => panic!("invalid timezone should fail"), + Err(error) => error, + }; match err { AppError::Internal(msg) => assert!(msg.contains("Invalid timezone"), "{msg}"), other => panic!("expected Internal, got {other:?}"), @@ -269,6 +1387,9 @@ mod tests { scheduler.add_job(&task, state).await.unwrap(); // The job must now be tracked in the internal job map. assert!(scheduler.job_map.contains_key("t-ok")); + let current_job_id = *scheduler.job_map.get("t-ok").unwrap(); + assert!(scheduler.is_current_job("t-ok", current_job_id)); + assert!(!scheduler.is_current_job("t-ok", uuid::Uuid::nil())); } // ---- remove_job ---------------------------------------------------- @@ -282,23 +1403,31 @@ mod tests { } #[tokio::test] - async fn test_remove_job_cancels_active_run_and_clears_map() { + async fn test_remove_job_cancels_active_run_and_clears_job_map() { let (state, _db, _dir) = build_test_state().await; let scheduler = TaskScheduler::new("UTC").await.unwrap(); let task = make_task("t-remove", Some(FAR_FUTURE_CRON), true); scheduler.add_job(&task, state).await.unwrap(); assert!(scheduler.job_map.contains_key("t-remove")); - // Simulate an in-flight run to exercise the cancel_active_run path. - let token = CancellationToken::new(); - let token_clone = token.clone(); + // Simulate the executor guard draining after it observes cancellation. + let active_run = test_active_run("run-x"); + let cancellation = active_run.cancellation.clone(); + let completed = active_run.completed.clone(); + let active_runs = Arc::clone(&scheduler.active_runs); scheduler .active_runs - .insert("t-remove".to_string(), ("run-x".to_string(), token)); + .insert("t-remove".to_string(), active_run); + let cancelled = cancellation.clone(); + tokio::spawn(async move { + cancellation.cancelled().await; + clear_active_run_if_current(&active_runs, "t-remove", "run-x"); + completed.cancel(); + }); scheduler.remove_job("t-remove").await.unwrap(); - assert!(token_clone.is_cancelled(), "active run should be cancelled"); + assert!(cancelled.is_cancelled(), "active run should be cancelled"); assert!(!scheduler.is_running("t-remove")); assert!(!scheduler.job_map.contains_key("t-remove")); } @@ -356,25 +1485,628 @@ mod tests { } } - // ---- disable_job --------------------------------------------------- + #[test] + fn test_validate_cron_uses_scheduler_syntax() { + assert!(validate_cron(FAR_FUTURE_CRON).is_ok()); + let error = validate_cron("0 0 0 1 1 * 2027").unwrap_err(); + assert!(matches!(error, AppError::Validation(_))); + } + + #[tokio::test] + async fn test_next_run_uses_scheduler_weekday_semantics() { + let scheduler = TaskScheduler::new("UTC").await.unwrap(); + let next = scheduler + .next_run_at("0 0 0 * * 1") + .expect("weekday schedule should have a next run"); + assert_eq!(next.weekday(), Weekday::Mon); + } #[tokio::test] - async fn test_disable_job_removes_registered_job() { + async fn test_create_scheduled_task_owns_row_and_job() { let (state, _db, _dir) = build_test_state().await; - let scheduler = TaskScheduler::new("UTC").await.unwrap(); - let task = make_task("t-disable", Some(FAR_FUTURE_CRON), true); - scheduler.add_job(&task, state).await.unwrap(); - assert!(scheduler.job_map.contains_key("t-disable")); + let created = create_task(&state, create_scheduled_request(), "admin", "127.0.0.1") + .await + .unwrap(); - scheduler.disable_job("t-disable").await.unwrap(); - assert!(!scheduler.job_map.contains_key("t-disable")); + assert!(created.next_run_at.is_some()); + assert!(state.task_scheduler.job_map.contains_key(&created.id)); + assert!( + task::Entity::find_by_id(&created.id) + .one(&state.db) + .await + .unwrap() + .is_some() + ); } #[tokio::test] - async fn test_disable_job_unknown_is_ok() { - // Disabling an unknown task delegates to remove_job and must succeed. - let scheduler = TaskScheduler::new("UTC").await.unwrap(); - scheduler.disable_job("nope").await.unwrap(); - assert!(!scheduler.is_running("nope")); + async fn test_update_scheduled_task_disables_and_restores_job() { + let (state, _db, _dir) = build_test_state().await; + let created = create_task(&state, create_scheduled_request(), "admin", "127.0.0.1") + .await + .unwrap(); + let update = |enabled| UpdateTaskRequest { + name: None, + command: None, + server_ids: None, + cron_expression: None, + enabled: Some(enabled), + timeout: None, + retry_count: None, + retry_interval: None, + }; + + let disabled = update_task(&state, &created.id, update(false)) + .await + .unwrap(); + assert!(!disabled.enabled); + assert!(!state.task_scheduler.job_map.contains_key(&created.id)); + + let enabled = update_task(&state, &created.id, update(true)) + .await + .unwrap(); + assert!(enabled.enabled); + assert!(enabled.next_run_at.is_some()); + assert!(state.task_scheduler.job_map.contains_key(&created.id)); + } + + #[tokio::test] + async fn test_delete_task_removes_row_results_job_and_active_run() { + let (state, _db, _dir) = build_test_state().await; + let created = create_task(&state, create_scheduled_request(), "admin", "127.0.0.1") + .await + .unwrap(); + write_synthetic_result( + &state.db, + &created.id, + "run-delete", + "server-1", + -2, + "denied", + ) + .await + .unwrap(); + let active_run = test_active_run("run-delete"); + let cancellation = active_run.cancellation.clone(); + let completed = active_run.completed.clone(); + let active_runs = Arc::clone(&state.task_scheduler.active_runs); + let db = state.db.clone(); + let task_id = created.id.clone(); + state + .task_scheduler + .active_runs + .insert(created.id.clone(), active_run); + let cancelled = cancellation.clone(); + tokio::spawn(async move { + cancellation.cancelled().await; + write_synthetic_result(&db, &task_id, "run-delete", "server-1", -3, "late result") + .await + .unwrap(); + clear_active_run_if_current(&active_runs, &task_id, "run-delete"); + completed.cancel(); + }); + + delete_task(&state, &created.id).await.unwrap(); + + assert!(cancelled.is_cancelled()); + assert!(!state.task_scheduler.job_map.contains_key(&created.id)); + assert!(!state.task_scheduler.active_runs.contains_key(&created.id)); + assert_eq!(result_count(&state.db, &created.id).await, 0); + assert!( + task::Entity::find_by_id(&created.id) + .one(&state.db) + .await + .unwrap() + .is_none() + ); + } + + #[tokio::test] + async fn test_server_cleanup_updates_or_removes_tasks_and_scheduler_jobs() { + let (state, _db, _dir) = build_test_state().await; + let mut deleted_request = create_scheduled_request(); + deleted_request.server_ids = vec!["orphan".to_string()]; + let deleted = create_task(&state, deleted_request, "admin", "127.0.0.1") + .await + .unwrap(); + let mut retained_request = create_scheduled_request(); + retained_request.server_ids = vec!["orphan".to_string(), "keep".to_string()]; + let retained = create_task(&state, retained_request, "admin", "127.0.0.1") + .await + .unwrap(); + + let active_run = test_active_run("cleanup-run"); + let cancellation = active_run.cancellation.clone(); + let completed = active_run.completed.clone(); + let active_runs = Arc::clone(&state.task_scheduler.active_runs); + let db = state.db.clone(); + let retained_task_id = retained.id.clone(); + state + .task_scheduler + .active_runs + .insert(retained.id.clone(), active_run); + let cancelled = cancellation.clone(); + tokio::spawn(async move { + cancellation.cancelled().await; + write_synthetic_result( + &db, + &retained_task_id, + "cleanup-run", + "orphan", + -3, + "late orphan result", + ) + .await + .unwrap(); + clear_active_run_if_current(&active_runs, &retained_task_id, "cleanup-run"); + completed.cancel(); + }); + + let mut cleanup = begin_server_cleanup(&state).await; + let transaction = state.db.begin().await.unwrap(); + cleanup + .remove_server_references(&transaction, &["orphan".to_string()]) + .await + .unwrap(); + transaction.commit().await.unwrap(); + cleanup.apply_after_commit(&state).await; + + assert!(cancelled.is_cancelled()); + assert!( + task::Entity::find_by_id(&deleted.id) + .one(&state.db) + .await + .unwrap() + .is_none() + ); + assert!(!state.task_scheduler.job_map.contains_key(&deleted.id)); + + let retained = task::Entity::find_by_id(&retained.id) + .one(&state.db) + .await + .unwrap() + .expect("task with a remaining server should survive"); + assert_eq!( + serde_json::from_str::>(&retained.server_ids_json).unwrap(), + vec!["keep"] + ); + assert!(state.task_scheduler.job_map.contains_key(&retained.id)); + assert_eq!( + task_result::Entity::find() + .filter(task_result::Column::ServerId.eq("orphan")) + .count(&state.db) + .await + .unwrap(), + 0 + ); + } + + #[tokio::test] + async fn test_restore_and_start_registers_only_enabled_scheduled_tasks() { + let (state, _db, _dir) = build_test_state().await; + seed_task(&state.db, "enabled-task", &["server-1"]).await; + seed_task(&state.db, "disabled-task", &["server-1"]).await; + task::Entity::update_many() + .filter(task::Column::Id.eq("disabled-task")) + .col_expr(task::Column::Enabled, Expr::value(false)) + .exec(&state.db) + .await + .unwrap(); + + restore_and_start(state.clone()).await; + + assert!(state.task_scheduler.job_map.contains_key("enabled-task")); + assert!(!state.task_scheduler.job_map.contains_key("disabled-task")); + let enabled = task::Entity::find_by_id("enabled-task") + .one(&state.db) + .await + .unwrap() + .unwrap(); + assert!(enabled.next_run_at.is_some()); + } + + #[tokio::test] + async fn test_run_now_reports_overlap_as_conflict() { + let (state, _db, _dir) = build_test_state().await; + let created = create_task(&state, create_scheduled_request(), "admin", "127.0.0.1") + .await + .unwrap(); + state + .task_scheduler + .active_runs + .insert(created.id.clone(), test_active_run("existing-run")); + + let error = run_now(&state, &created.id, None).await.unwrap_err(); + assert!(matches!(error, AppError::Conflict(_))); + } + + #[test] + fn test_correlation_id_format() { + let cid = build_correlation_id("task-1", "run-abc", "srv-1", 1); + assert_eq!(cid, "task-1:run-abc:srv-1:1"); + } + + #[test] + fn test_correlation_id_uniqueness() { + let a = build_correlation_id("t1", "r1", "s1", 1); + let b = build_correlation_id("t1", "r1", "s2", 1); + let c = build_correlation_id("t1", "r1", "s1", 2); + assert_ne!(a, b); + assert_ne!(a, c); + } + + // build_correlation_id must not collapse empty segments — the colon-joined + // shape is preserved even when every component is the empty string. + #[test] + fn test_correlation_id_empty_segments_keep_delimiters() { + assert_eq!(build_correlation_id("", "", "", 0), ":::0"); + } + + // Negative and large attempt numbers are formatted verbatim (no clamping). + #[test] + fn test_correlation_id_attempt_boundaries() { + assert_eq!(build_correlation_id("t", "r", "s", -1), "t:r:s:-1"); + assert_eq!( + build_correlation_id("t", "r", "s", i32::MAX), + format!("t:r:s:{}", i32::MAX) + ); + } + + // ---- Helpers ------------------------------------------------------- + + /// Insert a server row with the given id and persisted capability mirror. + /// The server is NOT registered with the agent_manager, so it is offline. + async fn seed_server(db: &DatabaseConnection, id: &str, capabilities: i32) { + use crate::entity::server; + let now = chrono::Utc::now(); + server::ActiveModel { + id: Set(id.to_string()), + token_hash: Set(Some("hash".to_string())), + token_prefix: Set(Some("sb_pref".to_string())), + name: Set(format!("server-{id}")), + cpu_name: Set(None), + cpu_cores: Set(None), + cpu_arch: Set(None), + os: Set(None), + kernel_version: Set(None), + mem_total: Set(None), + swap_total: Set(None), + disk_total: Set(None), + ipv4: Set(None), + ipv6: Set(None), + region: Set(None), + country_code: Set(None), + geo_manual: Set(false), + virtualization: Set(None), + agent_version: Set(None), + group_id: Set(None), + weight: Set(0), + hidden: Set(false), + remark: Set(None), + public_remark: Set(None), + price: Set(None), + billing_cycle: Set(None), + currency: Set(None), + expired_at: Set(None), + traffic_limit: Set(None), + traffic_limit_type: Set(None), + billing_start_day: Set(None), + capabilities: Set(capabilities), + protocol_version: Set(1), + features: Set("[]".to_string()), + last_remote_addr: Set(None), + fingerprint: Set(None), + created_at: Set(now), + updated_at: Set(now), + } + .insert(db) + .await + .expect("seed server"); + } + + /// Insert a scheduled task row targeting `server_ids`. + async fn seed_task(db: &DatabaseConnection, id: &str, server_ids: &[&str]) { + let now = chrono::Utc::now(); + task::ActiveModel { + id: Set(id.to_string()), + command: Set("echo hi".to_string()), + server_ids_json: Set(serde_json::to_string(server_ids).unwrap()), + created_by: Set("tester".to_string()), + task_type: Set("scheduled".to_string()), + name: Set(Some("test task".to_string())), + cron_expression: Set(Some("0 0 0 1 1 *".to_string())), + enabled: Set(true), + timeout: Set(Some(5)), + retry_count: Set(0), + retry_interval: Set(1), + last_run_at: Set(None), + next_run_at: Set(None), + created_at: Set(now), + } + .insert(db) + .await + .expect("seed task"); + } + + /// Count persisted task_result rows for a given task id. + async fn result_count(db: &DatabaseConnection, task_id: &str) -> u64 { + task_result::Entity::find() + .filter(task_result::Column::TaskId.eq(task_id)) + .count(db) + .await + .expect("count task results") + } + + /// Poll until at least `expected` task_result rows exist for `task_id` (the + /// offline executor path runs inside a detached tokio task spawned by + /// `execute_scheduled_task`). + async fn wait_for_results(db: &DatabaseConnection, task_id: &str, expected: u64) -> u64 { + for _ in 0..100 { + let n = result_count(db, task_id).await; + if n >= expected { + return n; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + result_count(db, task_id).await + } + + // ---- write_result (pure DB write) --------------------------------- + + // write_result persists a row with every field set verbatim, including a + // negative synthetic exit code and the supplied run_id/attempt/timestamps. + #[tokio::test] + async fn test_write_result_persists_all_fields() { + let (db, _tmp) = crate::test_utils::setup_test_db().await; + let started = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + write_result( + &db, + "task-w", + "run-w", + "srv-w", + 2, + started, + -3, + "Server offline", + ) + .await + .expect("write_result should insert"); + + let row = task_result::Entity::find() + .filter(task_result::Column::TaskId.eq("task-w")) + .one(&db) + .await + .unwrap() + .expect("row present"); + assert_eq!(row.server_id, "srv-w"); + assert_eq!(row.output, "Server offline"); + assert_eq!(row.exit_code, -3); + assert_eq!(row.run_id.as_deref(), Some("run-w")); + assert_eq!(row.attempt, 2); + assert_eq!(row.started_at, Some(started)); + } + + // ---- write_synthetic_result -------------------------------------- + + // write_synthetic_result delegates to write_result with a fixed attempt of + // 1 and stamps started_at itself (so it is Some, not the caller's value). + #[tokio::test] + async fn test_write_synthetic_result_uses_attempt_one() { + let (db, _tmp) = crate::test_utils::setup_test_db().await; + write_synthetic_result(&db, "task-s", "run-s", "srv-s", -2, "capability denied") + .await + .expect("synthetic write should insert"); + + let row = task_result::Entity::find() + .filter(task_result::Column::TaskId.eq("task-s")) + .one(&db) + .await + .unwrap() + .expect("row present"); + assert_eq!(row.attempt, 1); + assert_eq!(row.exit_code, -2); + assert_eq!(row.output, "capability denied"); + assert!(row.started_at.is_some()); + } + + // ---- execute_for_server (offline agent) --------------------------- + + // With no connected agent, execute_for_server falls into the get_sender + // None branch and writes exactly one "Server offline" result (exit -3) when + // retries are exhausted (retry_count == 0 → single attempt). + #[tokio::test] + async fn test_execute_for_server_offline_writes_single_result() { + let (state, _db, _dir) = build_test_state().await; + let token = CancellationToken::new(); + execute_for_server( + &state, "task-off", "run-off", "srv-off", "echo hi", 1, 0, 1, token, + ) + .await + .unwrap(); + + assert_eq!(result_count(&state.db, "task-off").await, 1); + let row = task_result::Entity::find() + .filter(task_result::Column::TaskId.eq("task-off")) + .one(&state.db) + .await + .unwrap() + .expect("offline result present"); + assert_eq!(row.exit_code, -3); + assert_eq!(row.output, "Server offline"); + assert_eq!(row.attempt, 1); + } + + // An already-cancelled token short-circuits before the attempt loop body, + // so no result row is written at all. + #[tokio::test] + async fn test_execute_for_server_cancelled_writes_nothing() { + let (state, _db, _dir) = build_test_state().await; + let token = CancellationToken::new(); + token.cancel(); + execute_for_server( + &state, + "task-cancel", + "run-c", + "srv-c", + "echo hi", + 1, + 2, + 1, + token, + ) + .await + .unwrap(); + + assert_eq!(result_count(&state.db, "task-cancel").await, 0); + } + + // ---- execute_scheduled_task: missing task ------------------------- + + // A task id that is not in the DB returns a real not-found error and clears + // its claimed active_runs entry (no result rows written). + #[tokio::test] + async fn test_execute_scheduled_task_missing_task_returns_not_found() { + let (state, _db, _dir) = build_test_state().await; + let error = execute_scheduled_task(&state, "ghost-task", true, None) + .await + .unwrap_err(); + assert!(matches!(error, AppError::NotFound(_))); + assert!(!state.task_scheduler.is_running("ghost-task")); + assert_eq!(result_count(&state.db, "ghost-task").await, 0); + } + + // ---- execute_scheduled_task: overlap guard ------------------------ + + // When an active run already occupies the active_runs slot, a second + // trigger is skipped (returns false) without touching the existing entry. + #[tokio::test] + async fn test_execute_scheduled_task_overlap_is_skipped() { + let (state, _db, _dir) = build_test_state().await; + seed_task(&state.db, "task-overlap", &[]).await; + state + .task_scheduler + .active_runs + .insert("task-overlap".to_string(), test_active_run("prior-run")); + + let started = execute_scheduled_task(&state, "task-overlap", true, None) + .await + .unwrap(); + assert!(!started); + // The pre-existing run_id must remain untouched. + let entry = state + .task_scheduler + .active_runs + .get("task-overlap") + .expect("entry retained"); + assert_eq!(entry.run_id, "prior-run"); + } + + // ---- execute_scheduled_task: capability denied -------------------- + + // A target server whose mirror lacks CAP_EXEC produces a synthetic + // capability-denied result (exit -2) written synchronously before any + // server execution is spawned. + #[tokio::test] + async fn test_execute_scheduled_task_cap_exec_denied_writes_synthetic() { + let (state, _db, _dir) = build_test_state().await; + // CAP_DEFAULT intentionally excludes CAP_EXEC. + seed_server(&state.db, "srv-nocap", CAP_DEFAULT as i32).await; + seed_task(&state.db, "task-nocap", &["srv-nocap"]).await; + + let started = execute_scheduled_task(&state, "task-nocap", true, None) + .await + .unwrap(); + assert!(started); + + let row = task_result::Entity::find() + .filter(task_result::Column::TaskId.eq("task-nocap")) + .one(&state.db) + .await + .unwrap() + .expect("synthetic denied result present"); + assert_eq!(row.exit_code, -2); + assert_eq!(row.server_id, "srv-nocap"); + assert!(row.output.contains("Capability denied")); + } + + // With an audit context supplied, the capability-denied branch also records + // an `exec_denied` audit log entry alongside the synthetic result. + #[tokio::test] + async fn test_execute_scheduled_task_cap_denied_logs_audit() { + use crate::entity::audit_log; + let (state, _db, _dir) = build_test_state().await; + seed_server(&state.db, "srv-audit", CAP_DEFAULT as i32).await; + seed_task(&state.db, "task-audit", &["srv-audit"]).await; + + let ctx = ExecAuditContext { + user_id: "admin".to_string(), + ip: "127.0.0.1".to_string(), + }; + let started = execute_scheduled_task(&state, "task-audit", true, Some(ctx)) + .await + .unwrap(); + assert!(started); + + let denied = audit_log::Entity::find() + .filter(audit_log::Column::Action.eq("exec_denied")) + .count(&state.db) + .await + .unwrap(); + assert_eq!(denied, 1); + } + + // ---- execute_scheduled_task: offline agent dispatch --------------- + + // A server that HAS CAP_EXEC but is not connected reaches execute_for_server + // via the spawned join_set, which writes a "Server offline" result (exit + // -3). retry_count is forced to 0 here because skip_retry == true. + #[tokio::test] + async fn test_execute_scheduled_task_offline_agent_writes_offline_result() { + let (state, _db, _dir) = build_test_state().await; + seed_server(&state.db, "srv-online-cap", (CAP_DEFAULT | CAP_EXEC) as i32).await; + seed_task(&state.db, "task-dispatch", &["srv-online-cap"]).await; + + let started = execute_scheduled_task(&state, "task-dispatch", true, None) + .await + .unwrap(); + assert!(started); + + let n = wait_for_results(&state.db, "task-dispatch", 1).await; + assert_eq!(n, 1); + let row = task_result::Entity::find() + .filter(task_result::Column::TaskId.eq("task-dispatch")) + .one(&state.db) + .await + .unwrap() + .expect("offline dispatch result present"); + assert_eq!(row.exit_code, -3); + assert_eq!(row.output, "Server offline"); + } + + // An empty server list produces no work and no result rows, but the trigger + // is still considered "started" (returns true) and updates last_run_at. + #[tokio::test] + async fn test_execute_scheduled_task_empty_server_list_starts_with_no_results() { + let (state, _db, _dir) = build_test_state().await; + seed_task(&state.db, "task-empty", &[]).await; + + let started = execute_scheduled_task(&state, "task-empty", true, None) + .await + .unwrap(); + assert!(started); + + // Give any (non-existent) spawned work a chance to run, then assert none. + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + assert_eq!(result_count(&state.db, "task-empty").await, 0); + + let updated = task::Entity::find_by_id("task-empty") + .one(&state.db) + .await + .unwrap() + .expect("task present"); + assert!( + updated.last_run_at.is_some(), + "last_run_at should be stamped" + ); } } diff --git a/crates/server/src/task/mod.rs b/crates/server/src/task/mod.rs index c0ada32d0..5560e56c1 100644 --- a/crates/server/src/task/mod.rs +++ b/crates/server/src/task/mod.rs @@ -5,5 +5,4 @@ pub mod offline_checker; pub mod record_writer; pub mod service_monitor_checker; pub mod session_cleaner; -pub mod task_scheduler; pub mod upgrade_timeout; diff --git a/crates/server/src/task/task_scheduler.rs b/crates/server/src/task/task_scheduler.rs deleted file mode 100644 index 48b9f871b..000000000 --- a/crates/server/src/task/task_scheduler.rs +++ /dev/null @@ -1,808 +0,0 @@ -use std::str::FromStr; -use std::sync::Arc; - -use chrono::Utc; -use dashmap::DashMap; -use dashmap::mapref::entry::Entry; -use sea_orm::prelude::Expr; -use sea_orm::*; -use tokio::task::JoinSet; -use tokio_util::sync::CancellationToken; - -use crate::entity::{task, task_result}; -use crate::service::agent_manager::AgentRequestError; -use crate::service::audit::AuditService; -use crate::service::high_risk_audit::ExecAuditContext; -use crate::state::AppState; -use serverbee_common::constants::CAP_EXEC; -use serverbee_common::protocol::{AgentMessage, ServerMessage}; - -/// Build correlation ID: {task_id}:{run_id}:{server_id}:{attempt} -pub fn build_correlation_id(task_id: &str, run_id: &str, server_id: &str, attempt: i32) -> String { - format!("{task_id}:{run_id}:{server_id}:{attempt}") -} - -/// Called by cron trigger or manual /run endpoint. -/// `skip_retry`: if true, executes once without retry (used by manual trigger). -/// Returns true if execution was started, false if skipped (overlap). -pub async fn execute_scheduled_task( - state: &Arc, - task_id: &str, - skip_retry: bool, - audit_context: Option, -) -> bool { - let scheduler = &state.task_scheduler; - - // Step 0: Atomic overlap check using DashMap entry API - let run_id = uuid::Uuid::new_v4().to_string(); - let token = CancellationToken::new(); - match scheduler.active_runs.entry(task_id.to_string()) { - Entry::Occupied(_) => { - tracing::warn!("Task {task_id} still running, skipping trigger"); - return false; - } - Entry::Vacant(e) => { - e.insert((run_id.clone(), token.clone())); - } - } - - // Step 1: Load task from DB (entry already claimed, remove on failure) - let task_model = match task::Entity::find_by_id(task_id).one(&state.db).await { - Ok(Some(t)) => t, - Ok(None) => { - tracing::error!("Task {task_id} not found"); - scheduler.active_runs.remove(task_id); - return false; - } - Err(e) => { - tracing::error!("Failed to load task {task_id}: {e}"); - scheduler.active_runs.remove(task_id); - return false; - } - }; - - let server_ids: Vec = - serde_json::from_str(&task_model.server_ids_json).unwrap_or_default(); - let timeout_secs = task_model.timeout.unwrap_or(300).max(1) as u64; - let retry_count = if skip_retry { - 0 - } else { - task_model.retry_count.max(0) - }; - let retry_interval = task_model.retry_interval.max(1) as u64; - let command = task_model.command.clone(); - let audit_context_ref = audit_context.as_ref(); - if let Some(context) = audit_context.clone() { - state.exec_audit_contexts.insert(run_id.clone(), context); - } - - // Step 3: Update last_run_at and compute next_run_at (using configured timezone) - let tz: chrono_tz::Tz = scheduler.timezone().parse().unwrap_or(chrono_tz::UTC); - let next_run = task_model.cron_expression.as_deref().and_then(|cron_expr| { - cron::Schedule::from_str(cron_expr) - .ok() - .and_then(|s| s.upcoming(tz).next().map(|dt| dt.with_timezone(&Utc))) - }); - let mut update = task::Entity::update_many() - .filter(task::Column::Id.eq(task_id)) - .col_expr(task::Column::LastRunAt, Expr::value(Utc::now())); - if let Some(next) = next_run { - update = update.col_expr(task::Column::NextRunAt, Expr::value(next)); - } - let _ = update.exec(&state.db).await; - - // Step 4: Dispatch to each server - let mut join_set = JoinSet::new(); - - // Get capabilities for target servers - use crate::entity::server; - let target_servers = server::Entity::find() - .filter(server::Column::Id.is_in(server_ids.clone())) - .all(&state.db) - .await - .unwrap_or_default(); - let server_caps: std::collections::HashMap = target_servers - .iter() - .map(|s| (s.id.clone(), s.capabilities)) - .collect(); - - for sid in &server_ids { - let caps = server_caps.get(sid).copied().unwrap_or(0); - - // Check CAP_EXEC - if let Some(reason) = - state - .agent_manager - .capability_denied_reason(sid, caps as u32, CAP_EXEC) - { - let _ = write_synthetic_result( - &state.db, - task_id, - &run_id, - sid, - -2, - crate::router::api::task::exec_capability_denied_output(reason), - ) - .await; - if let Some(context) = audit_context_ref { - let detail = serde_json::json!({ - "server_id": sid, - "task_id": task_id, - "command": command, - "deny_reason": reason, - }) - .to_string(); - let _ = AuditService::log( - &state.db, - &context.user_id, - "exec_denied", - Some(&detail), - &context.ip, - ) - .await; - } - continue; - } - - if let Some(context) = audit_context_ref { - let detail = serde_json::json!({ - "server_id": sid, - "task_id": task_id, - "command": command, - "timeout": Some(timeout_secs as u32), - }) - .to_string(); - let _ = AuditService::log( - &state.db, - &context.user_id, - "exec_started", - Some(&detail), - &context.ip, - ) - .await; - } - - let state = state.clone(); - let task_id = task_id.to_string(); - let run_id = run_id.clone(); - let sid = sid.clone(); - let command = command.clone(); - let token = token.clone(); - - join_set.spawn(async move { - execute_for_server( - &state, - &task_id, - &run_id, - &sid, - &command, - timeout_secs, - retry_count, - retry_interval, - token, - ) - .await; - }); - } - - // Step 5: Wait for all to complete, then clear active_runs. - // Use Arc::clone so the guard operates on the *shared* DashMap, not a clone. - let task_id_owned = task_id.to_string(); - let active_runs = Arc::clone(&scheduler.active_runs); - let run_id_for_cleanup = run_id.clone(); - let state_for_cleanup = state.clone(); - tokio::spawn(async move { - struct ActiveRunGuard { - active_runs: Arc>, - task_id: String, - state: Arc, - run_id: String, - } - impl Drop for ActiveRunGuard { - fn drop(&mut self) { - self.active_runs.remove(&self.task_id); - self.state.exec_audit_contexts.remove(&self.run_id); - } - } - let _guard = ActiveRunGuard { - active_runs, - task_id: task_id_owned, - state: state_for_cleanup, - run_id: run_id_for_cleanup, - }; - while join_set.join_next().await.is_some() {} - }); - - true -} - -#[allow(clippy::too_many_arguments)] -async fn execute_for_server( - state: &Arc, - task_id: &str, - run_id: &str, - server_id: &str, - command: &str, - timeout_secs: u64, - retry_count: i32, - retry_interval: u64, - token: CancellationToken, -) { - let max_attempts = retry_count + 1; - let timeout_duration = std::time::Duration::from_secs(timeout_secs + 10); - - for attempt in 1..=max_attempts { - if token.is_cancelled() { - break; - } - - if attempt > 1 { - // Wait for retry interval, but abort early if cancelled - tokio::select! { - _ = tokio::time::sleep(std::time::Duration::from_secs(retry_interval)) => {} - _ = token.cancelled() => { break; } - } - } - - let correlation_id = build_correlation_id(task_id, run_id, server_id, attempt); - let started_at = Utc::now(); - - // Send Exec and await the correlated TaskResult, aborting if the task - // is cancelled (deleted/disabled) - let result = tokio::select! { - r = state.agent_manager.request_with_id( - server_id, - correlation_id, - timeout_duration, - |msg_id| ServerMessage::Exec { - task_id: msg_id, - command: command.to_string(), - timeout: Some(timeout_secs as u32), - }, - ) => r, - _ = token.cancelled() => { - break; - } - }; - - // Double-check cancellation before writing — select! may resolve both - // futures simultaneously and pick the result branch. - if token.is_cancelled() { - break; - } - - match result { - Ok(AgentMessage::TaskResult { result, .. }) => { - let _ = write_result( - &state.db, - task_id, - run_id, - server_id, - attempt, - started_at, - result.exit_code, - &result.output, - ) - .await; - if result.exit_code == 0 { - break; - } - // Non-zero: continue to retry if attempts remain - } - Err(AgentRequestError::Offline) => { - let _ = write_result( - &state.db, - task_id, - run_id, - server_id, - attempt, - started_at, - -3, - "Server offline", - ) - .await; - } - Err(AgentRequestError::SendFailed) => { - let _ = write_result( - &state.db, - task_id, - run_id, - server_id, - attempt, - started_at, - -3, - "Dispatch failed", - ) - .await; - } - _ => { - // Timeout or channel error - let _ = write_result( - &state.db, - task_id, - run_id, - server_id, - attempt, - started_at, - -4, - &format!("No response within {timeout_secs}s"), - ) - .await; - } - } - - if attempt == max_attempts { - break; - } - } -} - -#[allow(clippy::too_many_arguments)] -async fn write_result( - db: &DatabaseConnection, - task_id: &str, - run_id: &str, - server_id: &str, - attempt: i32, - started_at: chrono::DateTime, - exit_code: i32, - output: &str, -) -> Result<(), DbErr> { - task_result::ActiveModel { - id: NotSet, - task_id: Set(task_id.to_string()), - server_id: Set(server_id.to_string()), - output: Set(output.to_string()), - exit_code: Set(exit_code), - finished_at: Set(Utc::now()), - run_id: Set(Some(run_id.to_string())), - attempt: Set(attempt), - started_at: Set(Some(started_at)), - } - .insert(db) - .await?; - Ok(()) -} - -async fn write_synthetic_result( - db: &DatabaseConnection, - task_id: &str, - run_id: &str, - server_id: &str, - exit_code: i32, - output: &str, -) -> Result<(), DbErr> { - write_result( - db, - task_id, - run_id, - server_id, - 1, - Utc::now(), - exit_code, - output, - ) - .await -} - -/// Startup function: load all enabled scheduled tasks and register jobs. -pub async fn run(state: Arc) { - let tasks = task::Entity::find() - .filter(task::Column::TaskType.eq("scheduled")) - .filter(task::Column::Enabled.eq(true)) - .all(&state.db) - .await; - - match tasks { - Ok(tasks) => { - let tz: chrono_tz::Tz = state - .task_scheduler - .timezone() - .parse() - .unwrap_or(chrono_tz::UTC); - for t in &tasks { - if let Err(e) = state.task_scheduler.add_job(t, state.clone()).await { - tracing::error!("Failed to register scheduled task {}: {e}", t.id); - continue; - } - // Refresh next_run_at on startup so it reflects the current time - if let Some(cron_expr) = &t.cron_expression { - let next = cron::Schedule::from_str(cron_expr) - .ok() - .and_then(|s| s.upcoming(tz).next().map(|dt| dt.with_timezone(&Utc))); - let _ = task::Entity::update_many() - .filter(task::Column::Id.eq(&t.id)) - .col_expr(task::Column::NextRunAt, Expr::value(next)) - .exec(&state.db) - .await; - } - } - tracing::info!("Loaded {} scheduled tasks", tasks.len()); - } - Err(e) => { - tracing::error!("Failed to load scheduled tasks: {e}"); - } - } - - // Start the scheduler tick loop - if let Err(e) = state.task_scheduler.start().await { - tracing::error!("Failed to start task scheduler: {e}"); - } -} - -#[cfg(test)] -mod tests { - use super::*; - use chrono::TimeZone; - use serverbee_common::constants::CAP_DEFAULT; - - #[test] - fn test_correlation_id_format() { - let cid = build_correlation_id("task-1", "run-abc", "srv-1", 1); - assert_eq!(cid, "task-1:run-abc:srv-1:1"); - } - - #[test] - fn test_correlation_id_uniqueness() { - let a = build_correlation_id("t1", "r1", "s1", 1); - let b = build_correlation_id("t1", "r1", "s2", 1); - let c = build_correlation_id("t1", "r1", "s1", 2); - assert_ne!(a, b); - assert_ne!(a, c); - } - - // build_correlation_id must not collapse empty segments — the colon-joined - // shape is preserved even when every component is the empty string. - #[test] - fn test_correlation_id_empty_segments_keep_delimiters() { - assert_eq!(build_correlation_id("", "", "", 0), ":::0"); - } - - // Negative and large attempt numbers are formatted verbatim (no clamping). - #[test] - fn test_correlation_id_attempt_boundaries() { - assert_eq!(build_correlation_id("t", "r", "s", -1), "t:r:s:-1"); - assert_eq!( - build_correlation_id("t", "r", "s", i32::MAX), - format!("t:r:s:{}", i32::MAX) - ); - } - - // ---- Helpers ------------------------------------------------------- - - /// Build an `Arc` backed by a fresh migrated test DB. The two - /// returned `TempDir` guards must outlive the test: one owns the SQLite - /// file, the other backs `data_dir` so GeoIP/ASN/transfer paths never - /// touch the working directory. - async fn build_test_state() -> ( - Arc, - tempfile::TempDir, - tempfile::TempDir, - ) { - let (db, db_guard) = crate::test_utils::setup_test_db().await; - let data_dir = tempfile::TempDir::new().unwrap(); - let mut config = crate::config::AppConfig::default(); - config.server.data_dir = data_dir.path().to_str().unwrap().to_string(); - let state = AppState::new(db, config).await.unwrap(); - (state, db_guard, data_dir) - } - - /// Insert a server row with the given id and persisted capability mirror. - /// The server is NOT registered with the agent_manager, so it is offline. - async fn seed_server(db: &DatabaseConnection, id: &str, capabilities: i32) { - use crate::entity::server; - let now = chrono::Utc::now(); - server::ActiveModel { - id: Set(id.to_string()), - token_hash: Set(Some("hash".to_string())), - token_prefix: Set(Some("sb_pref".to_string())), - name: Set(format!("server-{id}")), - cpu_name: Set(None), - cpu_cores: Set(None), - cpu_arch: Set(None), - os: Set(None), - kernel_version: Set(None), - mem_total: Set(None), - swap_total: Set(None), - disk_total: Set(None), - ipv4: Set(None), - ipv6: Set(None), - region: Set(None), - country_code: Set(None), - geo_manual: Set(false), - virtualization: Set(None), - agent_version: Set(None), - group_id: Set(None), - weight: Set(0), - hidden: Set(false), - remark: Set(None), - public_remark: Set(None), - price: Set(None), - billing_cycle: Set(None), - currency: Set(None), - expired_at: Set(None), - traffic_limit: Set(None), - traffic_limit_type: Set(None), - billing_start_day: Set(None), - capabilities: Set(capabilities), - protocol_version: Set(1), - features: Set("[]".to_string()), - last_remote_addr: Set(None), - fingerprint: Set(None), - created_at: Set(now), - updated_at: Set(now), - } - .insert(db) - .await - .expect("seed server"); - } - - /// Insert a scheduled task row targeting `server_ids`. - async fn seed_task(db: &DatabaseConnection, id: &str, server_ids: &[&str]) { - let now = chrono::Utc::now(); - task::ActiveModel { - id: Set(id.to_string()), - command: Set("echo hi".to_string()), - server_ids_json: Set(serde_json::to_string(server_ids).unwrap()), - created_by: Set("tester".to_string()), - task_type: Set("scheduled".to_string()), - name: Set(Some("test task".to_string())), - cron_expression: Set(Some("0 0 0 1 1 *".to_string())), - enabled: Set(true), - timeout: Set(Some(5)), - retry_count: Set(0), - retry_interval: Set(1), - last_run_at: Set(None), - next_run_at: Set(None), - created_at: Set(now), - } - .insert(db) - .await - .expect("seed task"); - } - - /// Count persisted task_result rows for a given task id. - async fn result_count(db: &DatabaseConnection, task_id: &str) -> u64 { - task_result::Entity::find() - .filter(task_result::Column::TaskId.eq(task_id)) - .count(db) - .await - .expect("count task results") - } - - /// Poll until at least `expected` task_result rows exist for `task_id` (the - /// offline executor path runs inside a detached tokio task spawned by - /// `execute_scheduled_task`). - async fn wait_for_results(db: &DatabaseConnection, task_id: &str, expected: u64) -> u64 { - for _ in 0..100 { - let n = result_count(db, task_id).await; - if n >= expected { - return n; - } - tokio::time::sleep(std::time::Duration::from_millis(20)).await; - } - result_count(db, task_id).await - } - - // ---- write_result (pure DB write) --------------------------------- - - // write_result persists a row with every field set verbatim, including a - // negative synthetic exit code and the supplied run_id/attempt/timestamps. - #[tokio::test] - async fn test_write_result_persists_all_fields() { - let (db, _tmp) = crate::test_utils::setup_test_db().await; - let started = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); - write_result(&db, "task-w", "run-w", "srv-w", 2, started, -3, "Server offline") - .await - .expect("write_result should insert"); - - let row = task_result::Entity::find() - .filter(task_result::Column::TaskId.eq("task-w")) - .one(&db) - .await - .unwrap() - .expect("row present"); - assert_eq!(row.server_id, "srv-w"); - assert_eq!(row.output, "Server offline"); - assert_eq!(row.exit_code, -3); - assert_eq!(row.run_id.as_deref(), Some("run-w")); - assert_eq!(row.attempt, 2); - assert_eq!(row.started_at, Some(started)); - } - - // ---- write_synthetic_result -------------------------------------- - - // write_synthetic_result delegates to write_result with a fixed attempt of - // 1 and stamps started_at itself (so it is Some, not the caller's value). - #[tokio::test] - async fn test_write_synthetic_result_uses_attempt_one() { - let (db, _tmp) = crate::test_utils::setup_test_db().await; - write_synthetic_result(&db, "task-s", "run-s", "srv-s", -2, "capability denied") - .await - .expect("synthetic write should insert"); - - let row = task_result::Entity::find() - .filter(task_result::Column::TaskId.eq("task-s")) - .one(&db) - .await - .unwrap() - .expect("row present"); - assert_eq!(row.attempt, 1); - assert_eq!(row.exit_code, -2); - assert_eq!(row.output, "capability denied"); - assert!(row.started_at.is_some()); - } - - // ---- execute_for_server (offline agent) --------------------------- - - // With no connected agent, execute_for_server falls into the get_sender - // None branch and writes exactly one "Server offline" result (exit -3) when - // retries are exhausted (retry_count == 0 → single attempt). - #[tokio::test] - async fn test_execute_for_server_offline_writes_single_result() { - let (state, _db, _dir) = build_test_state().await; - let token = CancellationToken::new(); - execute_for_server(&state, "task-off", "run-off", "srv-off", "echo hi", 1, 0, 1, token) - .await; - - assert_eq!(result_count(&state.db, "task-off").await, 1); - let row = task_result::Entity::find() - .filter(task_result::Column::TaskId.eq("task-off")) - .one(&state.db) - .await - .unwrap() - .expect("offline result present"); - assert_eq!(row.exit_code, -3); - assert_eq!(row.output, "Server offline"); - assert_eq!(row.attempt, 1); - } - - // An already-cancelled token short-circuits before the attempt loop body, - // so no result row is written at all. - #[tokio::test] - async fn test_execute_for_server_cancelled_writes_nothing() { - let (state, _db, _dir) = build_test_state().await; - let token = CancellationToken::new(); - token.cancel(); - execute_for_server(&state, "task-cancel", "run-c", "srv-c", "echo hi", 1, 2, 1, token) - .await; - - assert_eq!(result_count(&state.db, "task-cancel").await, 0); - } - - // ---- execute_scheduled_task: missing task ------------------------- - - // A task id that is not in the DB is rejected: the function returns false - // and clears its claimed active_runs entry (no result rows written). - #[tokio::test] - async fn test_execute_scheduled_task_missing_task_returns_false() { - let (state, _db, _dir) = build_test_state().await; - let started = execute_scheduled_task(&state, "ghost-task", true, None).await; - assert!(!started); - assert!(!state.task_scheduler.is_running("ghost-task")); - assert_eq!(result_count(&state.db, "ghost-task").await, 0); - } - - // ---- execute_scheduled_task: overlap guard ------------------------ - - // When an active run already occupies the active_runs slot, a second - // trigger is skipped (returns false) without touching the existing entry. - #[tokio::test] - async fn test_execute_scheduled_task_overlap_is_skipped() { - let (state, _db, _dir) = build_test_state().await; - seed_task(&state.db, "task-overlap", &[]).await; - let token = CancellationToken::new(); - state - .task_scheduler - .active_runs - .insert("task-overlap".to_string(), ("prior-run".to_string(), token)); - - let started = execute_scheduled_task(&state, "task-overlap", true, None).await; - assert!(!started); - // The pre-existing run_id must remain untouched. - let entry = state - .task_scheduler - .active_runs - .get("task-overlap") - .expect("entry retained"); - assert_eq!(entry.value().0, "prior-run"); - } - - // ---- execute_scheduled_task: capability denied -------------------- - - // A target server whose mirror lacks CAP_EXEC produces a synthetic - // capability-denied result (exit -2) written synchronously before any - // server execution is spawned. - #[tokio::test] - async fn test_execute_scheduled_task_cap_exec_denied_writes_synthetic() { - let (state, _db, _dir) = build_test_state().await; - // CAP_DEFAULT intentionally excludes CAP_EXEC. - seed_server(&state.db, "srv-nocap", CAP_DEFAULT as i32).await; - seed_task(&state.db, "task-nocap", &["srv-nocap"]).await; - - let started = execute_scheduled_task(&state, "task-nocap", true, None).await; - assert!(started); - - let row = task_result::Entity::find() - .filter(task_result::Column::TaskId.eq("task-nocap")) - .one(&state.db) - .await - .unwrap() - .expect("synthetic denied result present"); - assert_eq!(row.exit_code, -2); - assert_eq!(row.server_id, "srv-nocap"); - assert!(row.output.contains("Capability denied")); - } - - // With an audit context supplied, the capability-denied branch also records - // an `exec_denied` audit log entry alongside the synthetic result. - #[tokio::test] - async fn test_execute_scheduled_task_cap_denied_logs_audit() { - use crate::entity::audit_log; - let (state, _db, _dir) = build_test_state().await; - seed_server(&state.db, "srv-audit", CAP_DEFAULT as i32).await; - seed_task(&state.db, "task-audit", &["srv-audit"]).await; - - let ctx = ExecAuditContext { - user_id: "admin".to_string(), - ip: "127.0.0.1".to_string(), - }; - let started = execute_scheduled_task(&state, "task-audit", true, Some(ctx)).await; - assert!(started); - - let denied = audit_log::Entity::find() - .filter(audit_log::Column::Action.eq("exec_denied")) - .count(&state.db) - .await - .unwrap(); - assert_eq!(denied, 1); - } - - // ---- execute_scheduled_task: offline agent dispatch --------------- - - // A server that HAS CAP_EXEC but is not connected reaches execute_for_server - // via the spawned join_set, which writes a "Server offline" result (exit - // -3). retry_count is forced to 0 here because skip_retry == true. - #[tokio::test] - async fn test_execute_scheduled_task_offline_agent_writes_offline_result() { - let (state, _db, _dir) = build_test_state().await; - seed_server(&state.db, "srv-online-cap", (CAP_DEFAULT | CAP_EXEC) as i32).await; - seed_task(&state.db, "task-dispatch", &["srv-online-cap"]).await; - - let started = execute_scheduled_task(&state, "task-dispatch", true, None).await; - assert!(started); - - let n = wait_for_results(&state.db, "task-dispatch", 1).await; - assert_eq!(n, 1); - let row = task_result::Entity::find() - .filter(task_result::Column::TaskId.eq("task-dispatch")) - .one(&state.db) - .await - .unwrap() - .expect("offline dispatch result present"); - assert_eq!(row.exit_code, -3); - assert_eq!(row.output, "Server offline"); - } - - // An empty server list produces no work and no result rows, but the trigger - // is still considered "started" (returns true) and updates last_run_at. - #[tokio::test] - async fn test_execute_scheduled_task_empty_server_list_starts_with_no_results() { - let (state, _db, _dir) = build_test_state().await; - seed_task(&state.db, "task-empty", &[]).await; - - let started = execute_scheduled_task(&state, "task-empty", true, None).await; - assert!(started); - - // Give any (non-existent) spawned work a chance to run, then assert none. - tokio::time::sleep(std::time::Duration::from_millis(50)).await; - assert_eq!(result_count(&state.db, "task-empty").await, 0); - - let updated = task::Entity::find_by_id("task-empty") - .one(&state.db) - .await - .unwrap() - .expect("task present"); - assert!(updated.last_run_at.is_some(), "last_run_at should be stamped"); - } -} From cc97ddfd89cbcee7065080f896f81487577bb4dc Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 13:11:24 +0800 Subject: [PATCH 06/44] refactor(web): centralize server catalog projections Own the live, REST-list, and detail cache topology behind a single catalog interface. Route WebSocket frames, REST reads, and server mutations through typed projection events so transport adapters no longer coordinate query keys themselves.\n\nPreserve runtime metrics across REST snapshots, treat full sync and incremental frames according to their protocol semantics, guard refresh ordering and cache lifetime, and keep list/detail membership consistent after removals. Add catalog-level regression coverage and migrate all readers to the private namespace. --- .../dashboard/dashboard-editor-view.tsx | 2 +- .../components/dashboard/dashboard-grid.tsx | 2 +- .../dashboard/module-widget-host.tsx | 2 +- .../dashboard/widget-config-dialog.tsx | 2 +- .../dashboard/widget-render-dependencies.ts | 2 +- .../dashboard/widget-renderer.test.tsx | 2 +- .../components/dashboard/widget-renderer.tsx | 2 +- .../dashboard/widgets/alert-list.tsx | 2 +- .../components/dashboard/widgets/disk-io.tsx | 2 +- .../dashboard/widgets/gauge.test.tsx | 2 +- .../components/dashboard/widgets/gauge.tsx | 2 +- .../dashboard/widgets/line-chart-widget.tsx | 2 +- .../dashboard/widgets/metric-card.test.tsx | 2 +- .../dashboard/widgets/metric-card.tsx | 2 +- .../dashboard/widgets/multi-line.tsx | 2 +- .../widgets/network-latency-widget.tsx | 2 +- .../widgets/network-overview-widget.tsx | 2 +- .../dashboard/widgets/network-quality.tsx | 2 +- .../dashboard/widgets/server-cards.tsx | 2 +- .../dashboard/widgets/server-map.tsx | 2 +- .../dashboard/widgets/stat-number.test.tsx | 2 +- .../dashboard/widgets/stat-number.tsx | 2 +- .../components/dashboard/widgets/top-n.tsx | 2 +- .../dashboard/widgets/traffic-bar.tsx | 2 +- .../widgets/uptime-timeline-widget.tsx | 2 +- .../components/firewall/add-block-drawer.tsx | 15 +- .../server/add-server-dialog.test.tsx | 19 +- .../components/server/add-server-dialog.tsx | 13 +- .../components/server/pending-action-menu.tsx | 4 +- .../server/recover-agent-dialog.test.tsx | 40 +- .../server/recover-agent-dialog.tsx | 36 +- .../server/regenerate-code-dialog.tsx | 33 +- .../server/server-card-action-menu.tsx | 2 +- .../web/src/components/server/server-card.tsx | 2 +- .../components/server/server-edit-dialog.tsx | 13 +- .../status/server-detail-content.tsx | 13 +- .../components/task/scheduled-task-dialog.tsx | 15 +- .../components/task/scheduled-task-list.tsx | 13 +- apps/web/src/hooks/use-api.ts | 9 +- apps/web/src/hooks/use-metric-series.test.ts | 2 +- apps/web/src/hooks/use-metric-series.ts | 2 +- .../src/hooks/use-realtime-metrics.test.tsx | 48 +- apps/web/src/hooks/use-realtime-metrics.ts | 15 +- apps/web/src/hooks/use-server-tags.test.tsx | 23 +- apps/web/src/hooks/use-server-tags.ts | 6 +- apps/web/src/hooks/use-servers-ws.test.ts | 255 +------ apps/web/src/hooks/use-servers-ws.ts | 315 +-------- apps/web/src/hooks/use-upgrade-job.ts | 3 +- apps/web/src/lib/dev-mock-servers.ts | 2 +- apps/web/src/lib/server-catalog.test.ts | 552 ++++++++++++++++ apps/web/src/lib/server-catalog.ts | 620 ++++++++++++++++++ apps/web/src/lib/widget-helpers.ts | 2 +- apps/web/src/routes/_authed/index.tsx | 11 +- apps/web/src/routes/_authed/ip-quality.tsx | 13 +- .../src/routes/_authed/security/$serverId.tsx | 13 +- .../web/src/routes/_authed/security/index.tsx | 13 +- .../src/routes/_authed/servers/$id-page.tsx | 10 +- .../servers/$serverId/docker/index.tsx | 22 +- .../servers/components/index-cells.tsx | 2 +- .../servers/components/server-columns.tsx | 2 +- .../components/servers-page-toolbar.tsx | 2 +- .../_authed/servers/index.cells.test.tsx | 2 +- apps/web/src/routes/_authed/servers/index.tsx | 23 +- .../src/routes/_authed/settings/alerts.tsx | 12 +- .../routes/_authed/settings/capabilities.tsx | 15 +- .../routes/_authed/settings/ping-tasks.tsx | 8 +- .../routes/_authed/settings/status-pages.tsx | 9 +- .../web/src/routes/_authed/settings/tasks.tsx | 13 +- .../widgets-runtime/runtime-bridge.test.ts | 16 +- .../web/src/widgets-runtime/runtime-bridge.ts | 21 +- 70 files changed, 1410 insertions(+), 919 deletions(-) create mode 100644 apps/web/src/lib/server-catalog.test.ts create mode 100644 apps/web/src/lib/server-catalog.ts diff --git a/apps/web/src/components/dashboard/dashboard-editor-view.tsx b/apps/web/src/components/dashboard/dashboard-editor-view.tsx index 839f7ca80..ddb0b9a79 100644 --- a/apps/web/src/components/dashboard/dashboard-editor-view.tsx +++ b/apps/web/src/components/dashboard/dashboard-editor-view.tsx @@ -4,7 +4,7 @@ import { useTranslation } from 'react-i18next' import { Button } from '@/components/ui/button' import type { WidgetInput } from '@/hooks/use-dashboard' import { useDashboardEditor } from '@/hooks/use-dashboard-editor' -import type { ServerMetrics } from '@/hooks/use-servers-ws' +import type { ServerMetrics } from '@/lib/server-catalog' import type { Dashboard, DashboardWithWidgets } from '@/lib/widget-types' import { DashboardGrid } from './dashboard-grid' import { DashboardSwitcher } from './dashboard-switcher' diff --git a/apps/web/src/components/dashboard/dashboard-grid.tsx b/apps/web/src/components/dashboard/dashboard-grid.tsx index 9d59f2114..11f6ded77 100644 --- a/apps/web/src/components/dashboard/dashboard-grid.tsx +++ b/apps/web/src/components/dashboard/dashboard-grid.tsx @@ -29,7 +29,7 @@ import { import 'react-grid-layout/css/styles.css' import { useTranslation } from 'react-i18next' import { Button } from '@/components/ui/button' -import type { ServerMetrics } from '@/hooks/use-servers-ws' +import type { ServerMetrics } from '@/lib/server-catalog' import { cn } from '@/lib/utils' import type { DashboardWidget, SizingStrategy, WidgetTypeDefinition } from '@/lib/widget-types' import { WIDGET_TYPES } from '@/lib/widget-types' diff --git a/apps/web/src/components/dashboard/module-widget-host.tsx b/apps/web/src/components/dashboard/module-widget-host.tsx index 784c827af..2dd9e2071 100644 --- a/apps/web/src/components/dashboard/module-widget-host.tsx +++ b/apps/web/src/components/dashboard/module-widget-host.tsx @@ -2,7 +2,7 @@ import type { ActionsHelper } from '@serverbee/widget-sdk' import { useMemo } from 'react' import { useTranslation } from 'react-i18next' import { ScrollArea } from '@/components/ui/scroll-area' -import type { ServerMetrics } from '@/hooks/use-servers-ws' +import type { ServerMetrics } from '@/lib/server-catalog' import { parseConfig } from '@/lib/widget-helpers' import type { DashboardWidget } from '@/lib/widget-types' import { useWidgetRegistry } from '@/widgets-runtime/registry' diff --git a/apps/web/src/components/dashboard/widget-config-dialog.tsx b/apps/web/src/components/dashboard/widget-config-dialog.tsx index 544ed3583..b0acdebb3 100644 --- a/apps/web/src/components/dashboard/widget-config-dialog.tsx +++ b/apps/web/src/components/dashboard/widget-config-dialog.tsx @@ -11,7 +11,7 @@ import { Label } from '@/components/ui/label' import { ScrollArea } from '@/components/ui/scroll-area' import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from '@/components/ui/select' import { ToggleGroup, ToggleGroupItem } from '@/components/ui/toggle-group' -import type { ServerMetrics } from '@/hooks/use-servers-ws' +import type { ServerMetrics } from '@/lib/server-catalog' import { parseConfig } from '@/lib/widget-helpers' import type { AlertListConfig, diff --git a/apps/web/src/components/dashboard/widget-render-dependencies.ts b/apps/web/src/components/dashboard/widget-render-dependencies.ts index fc5c49a9b..3f29facc6 100644 --- a/apps/web/src/components/dashboard/widget-render-dependencies.ts +++ b/apps/web/src/components/dashboard/widget-render-dependencies.ts @@ -1,4 +1,4 @@ -import type { ServerMetrics } from '@/hooks/use-servers-ws' +import type { ServerMetrics } from '@/lib/server-catalog' import { parseConfig } from '@/lib/widget-helpers' import type { DashboardWidget } from '@/lib/widget-types' diff --git a/apps/web/src/components/dashboard/widget-renderer.test.tsx b/apps/web/src/components/dashboard/widget-renderer.test.tsx index 7331a57f4..2a88435e2 100644 --- a/apps/web/src/components/dashboard/widget-renderer.test.tsx +++ b/apps/web/src/components/dashboard/widget-renderer.test.tsx @@ -1,7 +1,7 @@ import type { WidgetManifest, WidgetModule } from '@serverbee/widget-sdk' import { act, render, screen } from '@testing-library/react' import { beforeEach, describe, expect, it, vi } from 'vitest' -import type { ServerMetrics } from '@/hooks/use-servers-ws' +import type { ServerMetrics } from '@/lib/server-catalog' import type { DashboardWidget } from '@/lib/widget-types' import { registryActions, useWidgetRegistry } from '@/widgets-runtime/registry' import { WidgetRenderer } from './widget-renderer' diff --git a/apps/web/src/components/dashboard/widget-renderer.tsx b/apps/web/src/components/dashboard/widget-renderer.tsx index 881a6a15e..c885712d2 100644 --- a/apps/web/src/components/dashboard/widget-renderer.tsx +++ b/apps/web/src/components/dashboard/widget-renderer.tsx @@ -1,5 +1,5 @@ import { Component, memo, type ReactNode, useMemo } from 'react' -import type { ServerMetrics } from '@/hooks/use-servers-ws' +import type { ServerMetrics } from '@/lib/server-catalog' import { parseConfig } from '@/lib/widget-helpers' import type { AlertListConfig, diff --git a/apps/web/src/components/dashboard/widgets/alert-list.tsx b/apps/web/src/components/dashboard/widgets/alert-list.tsx index 866e82676..e1117961f 100644 --- a/apps/web/src/components/dashboard/widgets/alert-list.tsx +++ b/apps/web/src/components/dashboard/widgets/alert-list.tsx @@ -2,8 +2,8 @@ import { useQuery } from '@tanstack/react-query' import { useMemo } from 'react' import { useTranslation } from 'react-i18next' import { ScrollArea } from '@/components/ui/scroll-area' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { api } from '@/lib/api-client' +import type { ServerMetrics } from '@/lib/server-catalog' import { filterByIds, formatRelativeTime } from '@/lib/widget-helpers' import type { AlertListConfig } from '@/lib/widget-types' diff --git a/apps/web/src/components/dashboard/widgets/disk-io.tsx b/apps/web/src/components/dashboard/widgets/disk-io.tsx index b3571d61a..9aa8f719b 100644 --- a/apps/web/src/components/dashboard/widgets/disk-io.tsx +++ b/apps/web/src/components/dashboard/widgets/disk-io.tsx @@ -11,8 +11,8 @@ import { import { CartesianGrid, Line, LineChart, XAxis, YAxis } from '@/components/ui/recharts-lazy' import { Skeleton } from '@/components/ui/skeleton' import { useServerRecords } from '@/hooks/use-api' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { buildMergedDiskIoSeries } from '@/lib/disk-io' +import type { ServerMetrics } from '@/lib/server-catalog' import { formatSpeed } from '@/lib/utils' import { formatChartTime } from '@/lib/widget-helpers' import type { DiskIoConfig } from '@/lib/widget-types' diff --git a/apps/web/src/components/dashboard/widgets/gauge.test.tsx b/apps/web/src/components/dashboard/widgets/gauge.test.tsx index 055430d68..6acfaf366 100644 --- a/apps/web/src/components/dashboard/widgets/gauge.test.tsx +++ b/apps/web/src/components/dashboard/widgets/gauge.test.tsx @@ -1,6 +1,6 @@ import { render, screen } from '@testing-library/react' import { describe, expect, it } from 'vitest' -import type { ServerMetrics } from '@/hooks/use-servers-ws' +import type { ServerMetrics } from '@/lib/server-catalog' import { GaugeWidget } from './gauge' function makeServer(id: string, overrides: Partial = {}): ServerMetrics { diff --git a/apps/web/src/components/dashboard/widgets/gauge.tsx b/apps/web/src/components/dashboard/widgets/gauge.tsx index b824c7832..87b89c254 100644 --- a/apps/web/src/components/dashboard/widgets/gauge.tsx +++ b/apps/web/src/components/dashboard/widgets/gauge.tsx @@ -1,7 +1,7 @@ import { Activity, Cpu, Gauge as GaugeIcon, HardDrive, MemoryStick, Network } from 'lucide-react' import { useId, useMemo } from 'react' import { useTranslation } from 'react-i18next' -import type { ServerMetrics } from '@/hooks/use-servers-ws' +import type { ServerMetrics } from '@/lib/server-catalog' import { extractLiveMetric, metricLabel } from '@/lib/widget-helpers' import type { GaugeConfig } from '@/lib/widget-types' diff --git a/apps/web/src/components/dashboard/widgets/line-chart-widget.tsx b/apps/web/src/components/dashboard/widgets/line-chart-widget.tsx index 6faf7839d..e38066a0c 100644 --- a/apps/web/src/components/dashboard/widgets/line-chart-widget.tsx +++ b/apps/web/src/components/dashboard/widgets/line-chart-widget.tsx @@ -4,7 +4,7 @@ import { type ChartConfig, ChartContainer, ChartTooltip, ChartTooltipContent } f import { Area, AreaChart, CartesianGrid, XAxis, YAxis } from '@/components/ui/recharts-lazy' import { Skeleton } from '@/components/ui/skeleton' import { useServerRecords } from '@/hooks/use-api' -import type { ServerMetrics } from '@/hooks/use-servers-ws' +import type { ServerMetrics } from '@/lib/server-catalog' import { formatBytes } from '@/lib/utils' import { extractRecordMetric, formatChartTime, isNetworkMetric, METRIC_UNITS, metricLabel } from '@/lib/widget-helpers' import type { LineChartConfig } from '@/lib/widget-types' diff --git a/apps/web/src/components/dashboard/widgets/metric-card.test.tsx b/apps/web/src/components/dashboard/widgets/metric-card.test.tsx index 6a938664d..99e13a627 100644 --- a/apps/web/src/components/dashboard/widgets/metric-card.test.tsx +++ b/apps/web/src/components/dashboard/widgets/metric-card.test.tsx @@ -2,7 +2,7 @@ import { QueryClient, QueryClientProvider } from '@tanstack/react-query' import { render, screen } from '@testing-library/react' import type { ReactNode } from 'react' import { beforeEach, describe, expect, it, vi } from 'vitest' -import type { ServerMetrics } from '@/hooks/use-servers-ws' +import type { ServerMetrics } from '@/lib/server-catalog' import { MetricCardWidget } from './metric-card' const translations: Record = { diff --git a/apps/web/src/components/dashboard/widgets/metric-card.tsx b/apps/web/src/components/dashboard/widgets/metric-card.tsx index cc1766751..3acc1821e 100644 --- a/apps/web/src/components/dashboard/widgets/metric-card.tsx +++ b/apps/web/src/components/dashboard/widgets/metric-card.tsx @@ -2,7 +2,7 @@ import { useMemo } from 'react' import { useTranslation } from 'react-i18next' import { useServerRecords } from '@/hooks/use-api' import { useMetricSeries } from '@/hooks/use-metric-series' -import type { ServerMetrics } from '@/hooks/use-servers-ws' +import type { ServerMetrics } from '@/lib/server-catalog' import { cn } from '@/lib/utils' import type { MetricCardConfig } from '@/lib/widget-types' import { METRIC_CARD_SPECS } from './metric-card/metric-card-config' diff --git a/apps/web/src/components/dashboard/widgets/multi-line.tsx b/apps/web/src/components/dashboard/widgets/multi-line.tsx index 9136eb905..1e776c4e9 100644 --- a/apps/web/src/components/dashboard/widgets/multi-line.tsx +++ b/apps/web/src/components/dashboard/widgets/multi-line.tsx @@ -11,9 +11,9 @@ import { } from '@/components/ui/chart' import { CartesianGrid, Line, LineChart, XAxis, YAxis } from '@/components/ui/recharts-lazy' import { Skeleton } from '@/components/ui/skeleton' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { api } from '@/lib/api-client' import type { ServerMetricRecord } from '@/lib/api-schema' +import type { ServerMetrics } from '@/lib/server-catalog' import { formatBytes } from '@/lib/utils' import { extractRecordMetric, diff --git a/apps/web/src/components/dashboard/widgets/network-latency-widget.tsx b/apps/web/src/components/dashboard/widgets/network-latency-widget.tsx index 674a9a881..dccedb4f8 100644 --- a/apps/web/src/components/dashboard/widgets/network-latency-widget.tsx +++ b/apps/web/src/components/dashboard/widgets/network-latency-widget.tsx @@ -4,8 +4,8 @@ import { LatencyChart } from '@/components/network/latency-chart' import { Skeleton } from '@/components/ui/skeleton' import { useNetworkServerSummary } from '@/hooks/use-network-api' import { useNetworkChartRecords } from '@/hooks/use-network-chart-records' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { CHART_COLORS } from '@/lib/chart-colors' +import type { ServerMetrics } from '@/lib/server-catalog' import type { NetworkLatencyConfig } from '@/lib/widget-types' interface NetworkLatencyWidgetProps { diff --git a/apps/web/src/components/dashboard/widgets/network-overview-widget.tsx b/apps/web/src/components/dashboard/widgets/network-overview-widget.tsx index 461f3dac1..baaa6a7e1 100644 --- a/apps/web/src/components/dashboard/widgets/network-overview-widget.tsx +++ b/apps/web/src/components/dashboard/widgets/network-overview-widget.tsx @@ -4,8 +4,8 @@ import { useTranslation } from 'react-i18next' import { ScrollArea } from '@/components/ui/scroll-area' import { Skeleton } from '@/components/ui/skeleton' import { useNetworkOverview } from '@/hooks/use-network-api' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { formatLatency, type NetworkServerSummary } from '@/lib/network-types' +import type { ServerMetrics } from '@/lib/server-catalog' import type { NetworkOverviewConfig } from '@/lib/widget-types' interface NetworkOverviewWidgetProps { diff --git a/apps/web/src/components/dashboard/widgets/network-quality.tsx b/apps/web/src/components/dashboard/widgets/network-quality.tsx index 93b7e1754..3e8930bd6 100644 --- a/apps/web/src/components/dashboard/widgets/network-quality.tsx +++ b/apps/web/src/components/dashboard/widgets/network-quality.tsx @@ -2,9 +2,9 @@ import { useTranslation } from 'react-i18next' import { ScrollArea } from '@/components/ui/scroll-area' import { Skeleton } from '@/components/ui/skeleton' import { useNetworkServerSummary } from '@/hooks/use-network-api' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { CHART_COLORS } from '@/lib/chart-colors' import { formatLatency, formatPacketLoss, getLossTextClassName } from '@/lib/network-types' +import type { ServerMetrics } from '@/lib/server-catalog' import type { NetworkQualityConfig } from '@/lib/widget-types' interface NetworkQualityWidgetProps { diff --git a/apps/web/src/components/dashboard/widgets/server-cards.tsx b/apps/web/src/components/dashboard/widgets/server-cards.tsx index fdc98ba79..9194c54d8 100644 --- a/apps/web/src/components/dashboard/widgets/server-cards.tsx +++ b/apps/web/src/components/dashboard/widgets/server-cards.tsx @@ -5,8 +5,8 @@ import { useTranslation } from 'react-i18next' import { DataTable } from '@/components/data-table/data-table' import { ServerCard } from '@/components/server/server-card' import { useCostOverview } from '@/hooks/use-cost' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { useTrafficOverview } from '@/hooks/use-traffic-overview' +import type { ServerMetrics } from '@/lib/server-catalog' import { filterByIds } from '@/lib/widget-helpers' import type { ServerCardsConfig } from '@/lib/widget-types' import { buildServerColumns } from '@/routes/_authed/servers/components/server-columns' diff --git a/apps/web/src/components/dashboard/widgets/server-map.tsx b/apps/web/src/components/dashboard/widgets/server-map.tsx index 796f31580..e5f60d5f8 100644 --- a/apps/web/src/components/dashboard/widgets/server-map.tsx +++ b/apps/web/src/components/dashboard/widgets/server-map.tsx @@ -5,8 +5,8 @@ import { useTranslation } from 'react-i18next' import { toast } from 'sonner' import { Button } from '@/components/ui/button' import { useAuth } from '@/hooks/use-auth' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { api } from '@/lib/api-client' +import type { ServerMetrics } from '@/lib/server-catalog' import { filterByIds } from '@/lib/widget-helpers' import type { ServerMapConfig } from '@/lib/widget-types' import { WORLD_MAP_PATHS } from '@/lib/world-map-paths' diff --git a/apps/web/src/components/dashboard/widgets/stat-number.test.tsx b/apps/web/src/components/dashboard/widgets/stat-number.test.tsx index b9fc64feb..9c226a0c3 100644 --- a/apps/web/src/components/dashboard/widgets/stat-number.test.tsx +++ b/apps/web/src/components/dashboard/widgets/stat-number.test.tsx @@ -1,6 +1,6 @@ import { render, screen } from '@testing-library/react' import { describe, expect, it, vi } from 'vitest' -import type { ServerMetrics } from '@/hooks/use-servers-ws' +import type { ServerMetrics } from '@/lib/server-catalog' import { StatNumberWidget } from './stat-number' const translations: Record = { diff --git a/apps/web/src/components/dashboard/widgets/stat-number.tsx b/apps/web/src/components/dashboard/widgets/stat-number.tsx index 56ef5765d..0abacb5b4 100644 --- a/apps/web/src/components/dashboard/widgets/stat-number.tsx +++ b/apps/web/src/components/dashboard/widgets/stat-number.tsx @@ -1,6 +1,6 @@ import { Activity, Cpu, MemoryStick, Server, Wifi } from 'lucide-react' import { useTranslation } from 'react-i18next' -import type { ServerMetrics } from '@/hooks/use-servers-ws' +import type { ServerMetrics } from '@/lib/server-catalog' import { cn, formatBytes } from '@/lib/utils' import type { StatNumberConfig } from '@/lib/widget-types' diff --git a/apps/web/src/components/dashboard/widgets/top-n.tsx b/apps/web/src/components/dashboard/widgets/top-n.tsx index 1e977325d..6ef18d90f 100644 --- a/apps/web/src/components/dashboard/widgets/top-n.tsx +++ b/apps/web/src/components/dashboard/widgets/top-n.tsx @@ -1,6 +1,6 @@ import { useMemo } from 'react' import { useTranslation } from 'react-i18next' -import type { ServerMetrics } from '@/hooks/use-servers-ws' +import type { ServerMetrics } from '@/lib/server-catalog' import { cn, formatBytes } from '@/lib/utils' import { extractLiveMetric, metricLabel } from '@/lib/widget-helpers' import type { TopNConfig } from '@/lib/widget-types' diff --git a/apps/web/src/components/dashboard/widgets/traffic-bar.tsx b/apps/web/src/components/dashboard/widgets/traffic-bar.tsx index 2d9fab9c1..f7cbd3e50 100644 --- a/apps/web/src/components/dashboard/widgets/traffic-bar.tsx +++ b/apps/web/src/components/dashboard/widgets/traffic-bar.tsx @@ -11,8 +11,8 @@ import { } from '@/components/ui/chart' import { Bar, BarChart, CartesianGrid, XAxis, YAxis } from '@/components/ui/recharts-lazy' import { Skeleton } from '@/components/ui/skeleton' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { api } from '@/lib/api-client' +import type { ServerMetrics } from '@/lib/server-catalog' import { formatBytes } from '@/lib/utils' import type { TrafficBarConfig } from '@/lib/widget-types' diff --git a/apps/web/src/components/dashboard/widgets/uptime-timeline-widget.tsx b/apps/web/src/components/dashboard/widgets/uptime-timeline-widget.tsx index 6d33e8797..a3c9e549b 100644 --- a/apps/web/src/components/dashboard/widgets/uptime-timeline-widget.tsx +++ b/apps/web/src/components/dashboard/widgets/uptime-timeline-widget.tsx @@ -3,9 +3,9 @@ import { useMemo } from 'react' import { useTranslation } from 'react-i18next' import { ScrollArea } from '@/components/ui/scroll-area' import { UptimeTimeline } from '@/components/uptime/uptime-timeline' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { api } from '@/lib/api-client' import type { UptimeDailyEntry } from '@/lib/api-schema' +import type { ServerMetrics } from '@/lib/server-catalog' import { computeAggregateUptime } from '@/lib/widget-helpers' import type { UptimeTimelineConfig } from '@/lib/widget-types' diff --git a/apps/web/src/components/firewall/add-block-drawer.tsx b/apps/web/src/components/firewall/add-block-drawer.tsx index fc9a48c84..095e73cbc 100644 --- a/apps/web/src/components/firewall/add-block-drawer.tsx +++ b/apps/web/src/components/firewall/add-block-drawer.tsx @@ -1,4 +1,3 @@ -import { useQuery } from '@tanstack/react-query' import { useEffect, useReducer } from 'react' import { useTranslation } from 'react-i18next' import { toast } from 'sonner' @@ -11,12 +10,8 @@ import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from '@ import { Sheet, SheetContent, SheetDescription, SheetHeader, SheetTitle } from '@/components/ui/sheet' import { Textarea } from '@/components/ui/textarea' import { useCreateBlock } from '@/hooks/use-firewall-blocks' -import { ApiError, api } from '@/lib/api-client' - -interface ServerLite { - id: string - name: string -} +import { ApiError } from '@/lib/api-client' +import { useServerList } from '@/lib/server-catalog' export interface AddBlockInitialValues { comment?: string @@ -110,11 +105,7 @@ export function AddBlockDrawer({ open, onOpenChange, initialValues }: Props) { dispatchForm({ type: 'reset', initialValues }) }, [open, initialValues]) - const { data: servers } = useQuery({ - queryKey: ['servers', 'lite'], - queryFn: () => api.get('/api/servers'), - enabled: open && form.coverType !== 'all' - }) + const { data: servers } = useServerList({ enabled: open && form.coverType !== 'all' }) const createMutation = useCreateBlock() diff --git a/apps/web/src/components/server/add-server-dialog.test.tsx b/apps/web/src/components/server/add-server-dialog.test.tsx index 2e9ef84d2..b1ab66e3e 100644 --- a/apps/web/src/components/server/add-server-dialog.test.tsx +++ b/apps/web/src/components/server/add-server-dialog.test.tsx @@ -3,8 +3,8 @@ import type { ReactNode } from 'react' import { beforeEach, describe, expect, it, vi } from 'vitest' const mockPost = vi.fn() -const mockGet = vi.fn() -const mockSetQueryData = vi.fn() +const mockRefreshServerCatalog = vi.hoisted(() => vi.fn()) +const mockQueryClient = {} vi.mock('react-i18next', () => ({ useTranslation: () => ({ @@ -34,7 +34,7 @@ vi.mock('@tanstack/react-query', () => ({ } }), useQuery: () => ({ data: [], isLoading: false }), - useQueryClient: () => ({ setQueryData: mockSetQueryData }) + useQueryClient: () => mockQueryClient })) vi.mock('sonner', () => ({ @@ -46,11 +46,14 @@ vi.mock('sonner', () => ({ vi.mock('@/lib/api-client', () => ({ api: { - get: (path: string) => mockGet(path), post: (path: string, body: unknown) => mockPost(path, body) } })) +vi.mock('@/lib/server-catalog', () => ({ + refreshServerCatalog: mockRefreshServerCatalog +})) + vi.mock('@/components/ui/button', () => ({ Button: ({ children, @@ -113,8 +116,7 @@ const { AddServerDialog } = await import('./add-server-dialog') describe('AddServerDialog', () => { beforeEach(() => { mockPost.mockReset() - mockGet.mockReset() - mockSetQueryData.mockReset() + mockRefreshServerCatalog.mockReset() }) it('POSTs to /api/servers with the form payload and transitions to the install-command view on success', async () => { @@ -127,7 +129,7 @@ describe('AddServerDialog', () => { expires_at: '2030-01-01T00:00:00Z' } }) - mockGet.mockResolvedValueOnce([]) + mockRefreshServerCatalog.mockResolvedValueOnce(undefined) render() @@ -150,8 +152,7 @@ describe('AddServerDialog', () => { }) expect(screen.getByText('add_server.shown_once_warning')).toBeInTheDocument() - await waitFor(() => expect(mockGet).toHaveBeenCalledWith('/api/servers')) - expect(mockSetQueryData).toHaveBeenCalledWith(['servers'], expect.any(Function)) + await waitFor(() => expect(mockRefreshServerCatalog).toHaveBeenCalledWith(mockQueryClient)) expect(screen.queryByRole('button', { name: 'add_server.generate' })).not.toBeInTheDocument() }) diff --git a/apps/web/src/components/server/add-server-dialog.tsx b/apps/web/src/components/server/add-server-dialog.tsx index 387a7511f..b0eb08d54 100644 --- a/apps/web/src/components/server/add-server-dialog.tsx +++ b/apps/web/src/components/server/add-server-dialog.tsx @@ -11,10 +11,10 @@ import { Dialog, DialogBody, DialogContent, DialogFooter, DialogHeader, DialogTi import { Input } from '@/components/ui/input' import { Popover, PopoverContent, PopoverTrigger } from '@/components/ui/popover' import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from '@/components/ui/select' -import { reconcileServersFromRest, type ServerMetrics } from '@/hooks/use-servers-ws' import { api } from '@/lib/api-client' -import type { CreateServerRequest, CreateServerResponse, ServerGroup, ServerResponse } from '@/lib/api-schema' +import type { CreateServerRequest, CreateServerResponse, ServerGroup } from '@/lib/api-schema' import { CAP_DEFAULT, CAPABILITIES, hasCap } from '@/lib/capabilities' +import { refreshServerCatalog } from '@/lib/server-catalog' import { cn } from '@/lib/utils' const DEFAULT_CAP_KEYS = CAPABILITIES.flatMap((c) => (hasCap(CAP_DEFAULT, c.bit) ? [c.key] : [])) @@ -617,15 +617,8 @@ export function AddServerDialog({ open, onClose }: { onClose: () => void; open: mutationFn: (body: CreateServerRequest) => api.post('/api/servers', body), onSuccess: async (data) => { dispatch({ type: 'setIssued', value: data }) - // ['servers'] is a WS-fed cache (queryFn: () => []); invalidating it - // would wipe the list. Refresh membership from REST and reconcile — - // existing rows keep their runtime metrics, the brand-new row is - // inserted as a stub until the next WS push fills it in. try { - const fresh = await api.get('/api/servers') - queryClient.setQueryData(['servers'], (prev) => - reconcileServersFromRest(prev, fresh as unknown as Array & { id: string }>) - ) + await refreshServerCatalog(queryClient) } catch { // Best-effort: the new row will surface on the next WS full_sync. } diff --git a/apps/web/src/components/server/pending-action-menu.tsx b/apps/web/src/components/server/pending-action-menu.tsx index 88cc923f6..c268d3941 100644 --- a/apps/web/src/components/server/pending-action-menu.tsx +++ b/apps/web/src/components/server/pending-action-menu.tsx @@ -15,8 +15,8 @@ import { } from '@/components/ui/alert-dialog' import { Button } from '@/components/ui/button' import { DropdownMenu, DropdownMenuContent, DropdownMenuItem, DropdownMenuTrigger } from '@/components/ui/dropdown-menu' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { api } from '@/lib/api-client' +import { projectServerCatalog } from '@/lib/server-catalog' import { RegenerateCodeDialog } from './regenerate-code-dialog' interface PendingActionMenuProps { @@ -33,7 +33,7 @@ export function PendingActionMenu({ serverId, serverName }: PendingActionMenuPro const deleteMutation = useMutation({ mutationFn: () => api.delete(`/api/servers/${serverId}`), onSuccess: () => { - queryClient.setQueryData(['servers'], (prev) => prev?.filter((s) => s.id !== serverId)) + projectServerCatalog(queryClient, { kind: 'servers_removed', serverIds: [serverId] }) toast.success(t('servers:card_pending.deleted')) setConfirmDeleteOpen(false) }, diff --git a/apps/web/src/components/server/recover-agent-dialog.test.tsx b/apps/web/src/components/server/recover-agent-dialog.test.tsx index 475ee8dee..9e2d2b66f 100644 --- a/apps/web/src/components/server/recover-agent-dialog.test.tsx +++ b/apps/web/src/components/server/recover-agent-dialog.test.tsx @@ -6,7 +6,8 @@ import { CAP_DEFAULT } from '@/lib/capabilities' const mockPost = vi.fn() const mockDelete = vi.fn() -const mockSetQueryData = vi.fn() +const mockProjectServerCatalog = vi.hoisted(() => vi.fn()) +const mockQueryClient = {} vi.mock('react-i18next', () => ({ useTranslation: () => ({ @@ -35,7 +36,7 @@ vi.mock('@tanstack/react-query', () => ({ } } }), - useQueryClient: () => ({ setQueryData: mockSetQueryData }) + useQueryClient: () => mockQueryClient })) vi.mock('sonner', () => ({ @@ -62,6 +63,10 @@ vi.mock('@/lib/api-client', () => ({ } })) +vi.mock('@/lib/server-catalog', () => ({ + projectServerCatalog: mockProjectServerCatalog +})) + vi.mock('@/components/ui/button', () => ({ Button: ({ children, @@ -131,7 +136,7 @@ describe('RecoverAgentDialog', () => { beforeEach(() => { mockPost.mockReset() mockDelete.mockReset() - mockSetQueryData.mockReset() + mockProjectServerCatalog.mockReset() }) it('renders the server name in a read-only header (not as an input)', () => { @@ -179,17 +184,11 @@ describe('RecoverAgentDialog', () => { expect(screen.getByText('plaintext-recover-code')).toBeInTheDocument() }) expect(screen.getByText('add_server.shown_once_warning')).toBeInTheDocument() - expect(mockSetQueryData).toHaveBeenCalledWith(['servers'], expect.any(Function)) - // The updater should patch outstanding_enrollment and (since revoke_immediately - // defaulted to true) flip has_token / online to false. - const updater = mockSetQueryData.mock.calls[0][1] as ( - prev: Record[] | undefined - ) => Record[] | undefined - const patched = updater([{ id: 'srv-42', has_token: true, online: true, outstanding_enrollment: null }]) - expect(patched?.[0]).toMatchObject({ - has_token: false, - online: false, - outstanding_enrollment: { id: 'enr-9', code_prefix: 'plaint' } + expect(mockProjectServerCatalog).toHaveBeenCalledWith(mockQueryClient, { + kind: 'enrollment_changed', + serverId: 'srv-42', + outstandingEnrollment: expect.objectContaining({ id: 'enr-9', code_prefix: 'plaint' }), + tokenRevoked: true }) }) @@ -244,13 +243,12 @@ describe('RecoverAgentDialog', () => { await waitFor(() => expect(mockDelete).toHaveBeenCalledTimes(1)) expect(mockDelete).toHaveBeenCalledWith('/api/agent/enrollments/enr-out-1') await waitFor(() => { - expect(mockSetQueryData).toHaveBeenCalledWith(['servers'], expect.any(Function)) + expect(mockProjectServerCatalog).toHaveBeenCalledWith(mockQueryClient, { + kind: 'enrollment_changed', + serverId: 'srv-42', + outstandingEnrollment: null, + tokenRevoked: false + }) }) - // The updater should clear outstanding_enrollment on the affected row. - const updater = mockSetQueryData.mock.calls[0][1] as ( - prev: Record[] | undefined - ) => Record[] | undefined - const patched = updater([{ id: 'srv-42', outstanding_enrollment: { id: 'enr-out-1', code_prefix: 'abc123' } }]) - expect(patched?.[0]).toMatchObject({ outstanding_enrollment: null }) }) }) diff --git a/apps/web/src/components/server/recover-agent-dialog.tsx b/apps/web/src/components/server/recover-agent-dialog.tsx index 9d348ac4e..37840ebe1 100644 --- a/apps/web/src/components/server/recover-agent-dialog.tsx +++ b/apps/web/src/components/server/recover-agent-dialog.tsx @@ -6,10 +6,10 @@ import { toast } from 'sonner' import { Button } from '@/components/ui/button' import { Checkbox } from '@/components/ui/checkbox' import { Dialog, DialogBody, DialogContent, DialogFooter, DialogHeader, DialogTitle } from '@/components/ui/dialog' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { ApiError, api } from '@/lib/api-client' import type { OutstandingEnrollmentSummary, RecoverRequest, RecoverResponse, ServerResponse } from '@/lib/api-schema' import { CAP_DEFAULT, CAPABILITIES, hasCap } from '@/lib/capabilities' +import { projectServerCatalog } from '@/lib/server-catalog' import { cn } from '@/lib/utils' const DEFAULT_CAP_KEYS = CAPABILITIES.flatMap((c) => (hasCap(CAP_DEFAULT, c.bit) ? [c.key] : [])) @@ -81,12 +81,12 @@ function OutstandingNotice({ enrollment, onClose, serverId }: OutstandingNoticeP mutationFn: () => api.delete(`/api/agent/enrollments/${enrollment.id}`), onSuccess: () => { toast.success(t('recover_agent.revoked')) - // ['servers'] is a WS-fed cache (queryFn: () => []) — invalidating it - // would wipe the visible list. Clear the row's outstanding_enrollment - // locally; WS will reconcile any drift. - queryClient.setQueryData(['servers'], (prev) => - prev?.map((s) => (s.id === serverId ? { ...s, outstanding_enrollment: null } : s)) - ) + projectServerCatalog(queryClient, { + kind: 'enrollment_changed', + serverId, + outstandingEnrollment: null, + tokenRevoked: false + }) }, onError: (err: unknown) => { const message = err instanceof ApiError || err instanceof Error ? err.message : t('recover_agent.revoke_failed') @@ -168,11 +168,6 @@ function RecoverAgentDialogContent({ mutationFn: (body: RecoverRequest) => api.post(`/api/servers/${server.id}/recover`, body), onSuccess: (data, variables) => { setIssued(data) - // Patch ['servers'] in place — invalidating it would wipe the WS-fed - // cache (queryFn: () => []) and unmount this dialog's host card. - // If `revoke_immediately`, the server row also returned to pending: its - // token_hash was cleared and the agent WS was kicked, so reflect - // has_token=false / online=false until the agent re-enrolls. const revoked = variables.revoke_immediately const newOutstanding = { id: data.enrollment.id, @@ -180,17 +175,12 @@ function RecoverAgentDialogContent({ expires_at: data.enrollment.expires_at, created_at: new Date().toISOString() } - queryClient.setQueryData(['servers'], (prev) => - prev?.map((s) => - s.id === server.id - ? { - ...s, - outstanding_enrollment: newOutstanding, - ...(revoked ? { has_token: false, online: false } : {}) - } - : s - ) - ) + projectServerCatalog(queryClient, { + kind: 'enrollment_changed', + serverId: server.id, + outstandingEnrollment: newOutstanding, + tokenRevoked: revoked + }) }, onError: (err: unknown) => { const message = err instanceof ApiError || err instanceof Error ? err.message : t('recover_agent.generate_failed') diff --git a/apps/web/src/components/server/regenerate-code-dialog.tsx b/apps/web/src/components/server/regenerate-code-dialog.tsx index 3eded0993..0d240caca 100644 --- a/apps/web/src/components/server/regenerate-code-dialog.tsx +++ b/apps/web/src/components/server/regenerate-code-dialog.tsx @@ -5,9 +5,9 @@ import { useTranslation } from 'react-i18next' import { toast } from 'sonner' import { Button } from '@/components/ui/button' import { Dialog, DialogBody, DialogContent, DialogFooter, DialogHeader, DialogTitle } from '@/components/ui/dialog' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { ApiError, api } from '@/lib/api-client' import type { RegenerateCodeRequest, RegenerateCodeResponse } from '@/lib/api-schema' +import { projectServerCatalog } from '@/lib/server-catalog' interface RegenerateCodeDialogProps { onOpenChange: (open: boolean) => void @@ -43,26 +43,17 @@ function RegenerateCodeDialogContent({ setIssued(data) setErrorMessage(null) toast.success(t('servers:card_pending.regenerated')) - // The ['servers'] key is a WS-fed cache whose queryFn returns []. Calling - // `invalidateQueries` here would re-run that queryFn and wipe the visible - // list (along with the ServerCard hosting this dialog). Patch the affected - // row's outstanding_enrollment in place instead — the next WS push will - // overwrite anything we got wrong. - queryClient.setQueryData(['servers'], (prev) => - prev?.map((s) => - s.id === serverId - ? { - ...s, - outstanding_enrollment: { - id: data.enrollment.id, - code_prefix: data.enrollment.code_prefix, - expires_at: data.enrollment.expires_at, - created_at: new Date().toISOString() - } - } - : s - ) - ) + projectServerCatalog(queryClient, { + kind: 'enrollment_changed', + serverId, + outstandingEnrollment: { + id: data.enrollment.id, + code_prefix: data.enrollment.code_prefix, + expires_at: data.enrollment.expires_at, + created_at: new Date().toISOString() + }, + tokenRevoked: false + }) }, onError: (err: unknown) => { const message = diff --git a/apps/web/src/components/server/server-card-action-menu.tsx b/apps/web/src/components/server/server-card-action-menu.tsx index b397ccd44..b6760222a 100644 --- a/apps/web/src/components/server/server-card-action-menu.tsx +++ b/apps/web/src/components/server/server-card-action-menu.tsx @@ -3,8 +3,8 @@ import { useState } from 'react' import { useTranslation } from 'react-i18next' import { Button } from '@/components/ui/button' import { DropdownMenu, DropdownMenuContent, DropdownMenuItem, DropdownMenuTrigger } from '@/components/ui/dropdown-menu' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { CAP_DEFAULT } from '@/lib/capabilities' +import type { ServerMetrics } from '@/lib/server-catalog' import { RecoverAgentDialog } from './recover-agent-dialog' import { ServerCardEditDialog } from './server-card-edit-dialog' diff --git a/apps/web/src/components/server/server-card.tsx b/apps/web/src/components/server/server-card.tsx index 6d42ce08d..c6c6b149a 100644 --- a/apps/web/src/components/server/server-card.tsx +++ b/apps/web/src/components/server/server-card.tsx @@ -4,11 +4,11 @@ import { useTranslation } from 'react-i18next' import { CompactMetric } from '@/components/server/compact-metric' import { MetricValue } from '@/components/server/metric-value' import { useNetworkRealtime } from '@/hooks/use-network-realtime' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import type { TrafficOverviewItem } from '@/hooks/use-traffic-overview' import type { ServerCostOverview } from '@/lib/api-schema' import { isLatencyFailure } from '@/lib/network-latency-constants' import { latencyColorClass, type NetworkServerSummary } from '@/lib/network-types' +import type { ServerMetrics } from '@/lib/server-catalog' import { computeTrafficQuota } from '@/lib/traffic' import { cn, formatBytes, formatUptime } from '@/lib/utils' import { useUpgradeJobsStore } from '@/stores/upgrade-jobs-store' diff --git a/apps/web/src/components/server/server-edit-dialog.tsx b/apps/web/src/components/server/server-edit-dialog.tsx index a28a87786..b8e00d1b9 100644 --- a/apps/web/src/components/server/server-edit-dialog.tsx +++ b/apps/web/src/components/server/server-edit-dialog.tsx @@ -21,10 +21,10 @@ import { Input } from '@/components/ui/input' import { Popover, PopoverContent, PopoverTrigger } from '@/components/ui/popover' import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from '@/components/ui/select' import { useServerTags, useUpdateServerTags } from '@/hooks/use-server-tags' -import { applyServerEdit, type ServerMetrics } from '@/hooks/use-servers-ws' import { api } from '@/lib/api-client' import type { ServerGroup, ServerResponse, UpdateServerInput } from '@/lib/api-schema' import { buildCountryOptions, type CountryOption } from '@/lib/country-codes' +import { projectServerCatalog } from '@/lib/server-catalog' import { cn, countryCodeToFlag } from '@/lib/utils' const TAG_SPLIT_RE = /[\s,]+/ @@ -409,16 +409,7 @@ function ServerEditDialogContent({ server, onClose }: { onClose: () => void; ser const mutation = useMutation({ mutationFn: (payload: UpdateServerInput) => api.put(`/api/servers/${server.id}`, payload), onSuccess: (data) => { - queryClient.setQueryData(['servers', server.id], data) - queryClient.setQueryData(['servers'], (prev) => - prev - ? applyServerEdit(prev, server.id, { - name: data.name, - group_id: data.group_id ?? null, - country_code: data.country_code ?? null - }) - : prev - ) + projectServerCatalog(queryClient, { kind: 'server_saved', server: data }) } }) diff --git a/apps/web/src/components/status/server-detail-content.tsx b/apps/web/src/components/status/server-detail-content.tsx index 02dcddbff..86f297799 100644 --- a/apps/web/src/components/status/server-detail-content.tsx +++ b/apps/web/src/components/status/server-detail-content.tsx @@ -16,7 +16,6 @@ import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs' import { UptimeTimeline } from '@/components/uptime/uptime-timeline' import { useServerRecords, useUptimeDaily } from '@/hooks/use-api' import { useRealtimeMetrics } from '@/hooks/use-realtime-metrics' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { api } from '@/lib/api-client' import type { PublicMetricsPoint, @@ -26,6 +25,7 @@ import type { UptimeDailyEntry } from '@/lib/api-schema' import { buildMergedDiskIoSeries, buildPerDiskIoSeries } from '@/lib/disk-io' +import { type ServerMetrics, useLiveServers } from '@/lib/server-catalog' import { cn, formatBytes } from '@/lib/utils' import { computeAggregateUptime } from '@/lib/widget-helpers' @@ -169,18 +169,11 @@ function useAdminGpuRecords(serverId: string, range: TimeRange, isAdminVariant: }) } -// Pulls the live-traffic strip data from the WS-driven `['servers']` cache. +// Pulls the live-traffic strip data from the WS-driven server catalog. // Public variant intentionally does not subscribe; the strip falls back to // the snapshot in `PublicServerDetail.metrics`. function useLiveServerMetrics(serverId: string, isAdminVariant: boolean) { - const { data: liveServers } = useQuery({ - queryKey: ['servers'], - queryFn: () => [], - staleTime: Number.POSITIVE_INFINITY, - refetchOnMount: false, - refetchOnWindowFocus: false, - enabled: isAdminVariant - }) + const { data: liveServers } = useLiveServers({ enabled: isAdminVariant }) return liveServers?.find((s) => s.id === serverId) } diff --git a/apps/web/src/components/task/scheduled-task-dialog.tsx b/apps/web/src/components/task/scheduled-task-dialog.tsx index d6071115b..a15344ae6 100644 --- a/apps/web/src/components/task/scheduled-task-dialog.tsx +++ b/apps/web/src/components/task/scheduled-task-dialog.tsx @@ -1,4 +1,3 @@ -import { useQuery } from '@tanstack/react-query' import { useReducer } from 'react' import { useTranslation } from 'react-i18next' import { toast } from 'sonner' @@ -13,8 +12,8 @@ import { useCreateScheduledTask, useUpdateScheduledTask } from '@/hooks/use-scheduled-tasks' -import { api } from '@/lib/api-client' import { CAP_EXEC, getEffectiveCapabilityEnabled } from '@/lib/capabilities' +import { useServerList } from '@/lib/server-catalog' const CRON_SPLIT_RE = /\s+/ @@ -23,13 +22,6 @@ interface Props { task?: ScheduledTask | null } -interface ServerInfo { - capabilities?: number - effective_capabilities?: number | null - id: string - name: string -} - interface ScheduledTaskFormState { command: string cronError: string @@ -99,10 +91,7 @@ export function ScheduledTaskDialog({ onClose, task }: Props) { const { t } = useTranslation(['settings', 'common']) const [state, dispatch] = useReducer(scheduledTaskFormReducer, task, scheduledTaskFormFromTask) - const { data: servers } = useQuery({ - queryKey: ['servers-list'], - queryFn: () => api.get('/api/servers') - }) + const { data: servers } = useServerList() const createMutation = useCreateScheduledTask() const updateMutation = useUpdateScheduledTask() diff --git a/apps/web/src/components/task/scheduled-task-list.tsx b/apps/web/src/components/task/scheduled-task-list.tsx index 48e38b73a..555088c0d 100644 --- a/apps/web/src/components/task/scheduled-task-list.tsx +++ b/apps/web/src/components/task/scheduled-task-list.tsx @@ -1,4 +1,3 @@ -import { useQuery } from '@tanstack/react-query' import { Calendar, ChevronDown, ChevronRight, Edit, Pause, Play, Plus, Trash2 } from 'lucide-react' import { useState } from 'react' import { useTranslation } from 'react-i18next' @@ -13,15 +12,10 @@ import { useScheduledTasks, useUpdateScheduledTask } from '@/hooks/use-scheduled-tasks' -import { api } from '@/lib/api-client' import { formatDateTime } from '@/lib/format' +import { useServerList } from '@/lib/server-catalog' import { ScheduledTaskDialog } from './scheduled-task-dialog' -interface ServerInfo { - id: string - name: string -} - function getExitCodeColor(code: number) { if (code === 0) { return 'bg-emerald-500/10 text-emerald-600 dark:text-emerald-400' @@ -47,10 +41,7 @@ export function ScheduledTaskList() { const runMutation = useRunScheduledTask() const updateMutation = useUpdateScheduledTask() - const { data: servers } = useQuery({ - queryKey: ['servers-list'], - queryFn: () => api.get('/api/servers') - }) + const { data: servers } = useServerList() const serverNameMap = new Map(servers?.map((s) => [s.id, s.name]) ?? []) diff --git a/apps/web/src/hooks/use-api.ts b/apps/web/src/hooks/use-api.ts index 933487b2c..6433397a8 100644 --- a/apps/web/src/hooks/use-api.ts +++ b/apps/web/src/hooks/use-api.ts @@ -1,15 +1,12 @@ import { useQuery } from '@tanstack/react-query' import { api } from '@/lib/api-client' -import type { ServerResponse, UptimeDailyEntry } from '@/lib/api-schema' +import type { UptimeDailyEntry } from '@/lib/api-schema' +import { useServerDetail } from '@/lib/server-catalog' type ServerRecord = import('@/lib/api-schema').ServerMetricRecord export function useServer(id: string) { - return useQuery({ - queryKey: ['servers', id], - queryFn: () => api.get(`/api/servers/${id}`), - enabled: !!id && id.length > 0 - }) + return useServerDetail(id) } export function useServerRecords(id: string, hours: number, interval: string, options?: { enabled?: boolean }) { diff --git a/apps/web/src/hooks/use-metric-series.test.ts b/apps/web/src/hooks/use-metric-series.test.ts index 4b2463264..6ce90910f 100644 --- a/apps/web/src/hooks/use-metric-series.test.ts +++ b/apps/web/src/hooks/use-metric-series.test.ts @@ -1,7 +1,7 @@ import { renderHook } from '@testing-library/react' import { describe, expect, it } from 'vitest' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import type { ServerMetricRecord } from '@/lib/api-schema' +import type { ServerMetrics } from '@/lib/server-catalog' import { useMetricSeries } from './use-metric-series' function record(time: string, cpu: number): ServerMetricRecord { diff --git a/apps/web/src/hooks/use-metric-series.ts b/apps/web/src/hooks/use-metric-series.ts index cfab59a3e..427a1102f 100644 --- a/apps/web/src/hooks/use-metric-series.ts +++ b/apps/web/src/hooks/use-metric-series.ts @@ -1,6 +1,6 @@ import { useMemo } from 'react' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import type { ServerMetricRecord } from '@/lib/api-schema' +import type { ServerMetrics } from '@/lib/server-catalog' import { extractLiveMetric, extractRecordMetric } from '@/lib/widget-helpers' export interface MetricSeriesPoint { diff --git a/apps/web/src/hooks/use-realtime-metrics.test.tsx b/apps/web/src/hooks/use-realtime-metrics.test.tsx index de7b756c7..afda09869 100644 --- a/apps/web/src/hooks/use-realtime-metrics.test.tsx +++ b/apps/web/src/hooks/use-realtime-metrics.test.tsx @@ -2,8 +2,8 @@ import { QueryClient, QueryClientProvider } from '@tanstack/react-query' import { act, renderHook } from '@testing-library/react' import type { ReactNode } from 'react' import { describe, expect, it } from 'vitest' +import { projectServerCatalog, type ServerMetrics } from '@/lib/server-catalog' import { toRealtimeDataPoint, useRealtimeMetrics } from './use-realtime-metrics' -import type { ServerMetrics } from './use-servers-ws' function makeMetrics(overrides: Partial = {}): ServerMetrics { return { @@ -40,6 +40,10 @@ function makeMetrics(overrides: Partial = {}): ServerMetrics { } } +function setLiveServers(queryClient: QueryClient, servers: ServerMetrics[]): void { + projectServerCatalog(queryClient, { kind: 'ws_full_sync', servers }) +} + describe('toRealtimeDataPoint', () => { it('converts with correct percentages', () => { const metrics = makeMetrics({ @@ -103,7 +107,7 @@ describe('useRealtimeMetrics', () => { it('seeds from existing cache on mount', () => { const { queryClient, Wrapper } = createTestEnv() - queryClient.setQueryData(['servers'], [makeMetrics({ id: 's1', cpu: 42, last_active: 1000 })]) + setLiveServers(queryClient, [makeMetrics({ id: 's1', cpu: 42, last_active: 1000 })]) const { result } = renderHook(() => useRealtimeMetrics('s1'), { wrapper: Wrapper }) @@ -122,7 +126,7 @@ describe('useRealtimeMetrics', () => { it('does not seed when server is offline', () => { const { queryClient, Wrapper } = createTestEnv() - queryClient.setQueryData(['servers'], [makeMetrics({ id: 's1', online: false })]) + setLiveServers(queryClient, [makeMetrics({ id: 's1', online: false })]) const { result } = renderHook(() => useRealtimeMetrics('s1'), { wrapper: Wrapper }) @@ -131,7 +135,7 @@ describe('useRealtimeMetrics', () => { it('does not seed when last_active is 0', () => { const { queryClient, Wrapper } = createTestEnv() - queryClient.setQueryData(['servers'], [makeMetrics({ id: 's1', last_active: 0 })]) + setLiveServers(queryClient, [makeMetrics({ id: 's1', last_active: 0 })]) const { result } = renderHook(() => useRealtimeMetrics('s1'), { wrapper: Wrapper }) @@ -140,13 +144,13 @@ describe('useRealtimeMetrics', () => { it('appends data point when last_active changes', () => { const { queryClient, Wrapper } = createTestEnv() - queryClient.setQueryData(['servers'], [makeMetrics({ id: 's1', cpu: 10, last_active: 1000 })]) + setLiveServers(queryClient, [makeMetrics({ id: 's1', cpu: 10, last_active: 1000 })]) const { result } = renderHook(() => useRealtimeMetrics('s1'), { wrapper: Wrapper }) expect(result.current).toHaveLength(1) act(() => { - queryClient.setQueryData(['servers'], [makeMetrics({ id: 's1', cpu: 20, last_active: 1003 })]) + setLiveServers(queryClient, [makeMetrics({ id: 's1', cpu: 20, last_active: 1003 })]) }) expect(result.current).toHaveLength(2) @@ -155,14 +159,14 @@ describe('useRealtimeMetrics', () => { it('deduplicates when last_active stays the same', () => { const { queryClient, Wrapper } = createTestEnv() - queryClient.setQueryData(['servers'], [makeMetrics({ id: 's1', cpu: 10, last_active: 1000 })]) + setLiveServers(queryClient, [makeMetrics({ id: 's1', cpu: 10, last_active: 1000 })]) const { result } = renderHook(() => useRealtimeMetrics('s1'), { wrapper: Wrapper }) expect(result.current).toHaveLength(1) // Same last_active, different cpu — should NOT append act(() => { - queryClient.setQueryData(['servers'], [makeMetrics({ id: 's1', cpu: 99, last_active: 1000 })]) + setLiveServers(queryClient, [makeMetrics({ id: 's1', cpu: 99, last_active: 1000 })]) }) expect(result.current).toHaveLength(1) @@ -171,20 +175,20 @@ describe('useRealtimeMetrics', () => { it('ignores updates for a different server', () => { const { queryClient, Wrapper } = createTestEnv() - queryClient.setQueryData( - ['servers'], - [makeMetrics({ id: 's1', last_active: 1000 }), makeMetrics({ id: 's2', last_active: 2000 })] - ) + setLiveServers(queryClient, [ + makeMetrics({ id: 's1', last_active: 1000 }), + makeMetrics({ id: 's2', last_active: 2000 }) + ]) const { result } = renderHook(() => useRealtimeMetrics('s1'), { wrapper: Wrapper }) expect(result.current).toHaveLength(1) // Update only s2 act(() => { - queryClient.setQueryData( - ['servers'], - [makeMetrics({ id: 's1', last_active: 1000 }), makeMetrics({ id: 's2', last_active: 2003 })] - ) + setLiveServers(queryClient, [ + makeMetrics({ id: 's1', last_active: 1000 }), + makeMetrics({ id: 's2', last_active: 2003 }) + ]) }) expect(result.current).toHaveLength(1) @@ -192,14 +196,14 @@ describe('useRealtimeMetrics', () => { it('trims buffer when exceeding threshold', () => { const { queryClient, Wrapper } = createTestEnv() - queryClient.setQueryData(['servers'], [makeMetrics({ id: 's1', last_active: 1 })]) + setLiveServers(queryClient, [makeMetrics({ id: 's1', last_active: 1 })]) const { result } = renderHook(() => useRealtimeMetrics('s1'), { wrapper: Wrapper }) // Push 260 updates (seed=1, then 259 more → 260 total, exceeds 250 threshold) act(() => { for (let i = 2; i <= 260; i++) { - queryClient.setQueryData(['servers'], [makeMetrics({ id: 's1', cpu: i, last_active: i })]) + setLiveServers(queryClient, [makeMetrics({ id: 's1', cpu: i, last_active: i })]) } }) @@ -213,10 +217,10 @@ describe('useRealtimeMetrics', () => { it('resets buffer when serverId changes', () => { const { queryClient, Wrapper } = createTestEnv() - queryClient.setQueryData( - ['servers'], - [makeMetrics({ id: 's1', cpu: 10, last_active: 1000 }), makeMetrics({ id: 's2', cpu: 20, last_active: 2000 })] - ) + setLiveServers(queryClient, [ + makeMetrics({ id: 's1', cpu: 10, last_active: 1000 }), + makeMetrics({ id: 's2', cpu: 20, last_active: 2000 }) + ]) const { result, rerender } = renderHook(({ id }) => useRealtimeMetrics(id), { initialProps: { id: 's1' }, diff --git a/apps/web/src/hooks/use-realtime-metrics.ts b/apps/web/src/hooks/use-realtime-metrics.ts index bffe5bfa3..e0f50ad1f 100644 --- a/apps/web/src/hooks/use-realtime-metrics.ts +++ b/apps/web/src/hooks/use-realtime-metrics.ts @@ -1,7 +1,7 @@ import type { QueryClient } from '@tanstack/react-query' import { useQueryClient } from '@tanstack/react-query' import { useEffect, useRef, useState } from 'react' -import type { ServerMetrics } from './use-servers-ws' +import { readLiveServers, type ServerMetrics, subscribeLiveServers } from '@/lib/server-catalog' const MAX_BUFFER_SIZE = 200 const TRIM_THRESHOLD = 250 @@ -62,7 +62,7 @@ export function toRealtimeDataPoint(metrics: ServerMetrics): RealtimeDataPoint { function createRealtimeBufferState(queryClient: QueryClient, serverId: string): RealtimeBufferState { const cache = getQueryClientBufferCache(queryClient) - const servers = queryClient.getQueryData(['servers']) + const servers = readLiveServers(queryClient) const server = servers?.find((s) => s.id === serverId) const cached = cache.get(serverId) @@ -113,7 +113,7 @@ export function useRealtimeMetrics(serverId: string): RealtimeDataPoint[] { return } const cache = getQueryClientBufferCache(queryClient) - const servers = queryClient.getQueryData(['servers']) + const servers = readLiveServers(queryClient) const server = servers?.find((s) => s.id === serverId) if (!server?.online || server.last_active <= 0) { cache.delete(serverId) @@ -134,13 +134,8 @@ export function useRealtimeMetrics(serverId: string): RealtimeDataPoint[] { }, RENDER_THROTTLE_MS) } - // Subscribe to cache updates - const unsubscribe = queryClient.getQueryCache().subscribe((event) => { - if (event.query.queryHash !== '["servers"]') { - return - } - - const data = event.query.state.data as ServerMetrics[] | undefined + const unsubscribe = subscribeLiveServers(queryClient, () => { + const data = readLiveServers(queryClient) if (!data) { return } diff --git a/apps/web/src/hooks/use-server-tags.test.tsx b/apps/web/src/hooks/use-server-tags.test.tsx index af3f67601..f029d955f 100644 --- a/apps/web/src/hooks/use-server-tags.test.tsx +++ b/apps/web/src/hooks/use-server-tags.test.tsx @@ -1,6 +1,7 @@ import { QueryClient, QueryClientProvider } from '@tanstack/react-query' import { renderHook, waitFor } from '@testing-library/react' import { afterEach, describe, expect, it, vi } from 'vitest' +import { projectServerCatalog, readLiveServers } from '@/lib/server-catalog' import { useServerTags, useUpdateServerTags } from './use-server-tags' function harness() { @@ -43,10 +44,28 @@ describe('useUpdateServerTags', () => { }) const { qc, Wrapper } = harness() qc.setQueryData(['server-tags', 'srv-1'], ['old']) - qc.setQueryData(['servers'], [{ id: 'srv-1', tags: ['old'] }]) + projectServerCatalog(qc, { + kind: 'rest_snapshot', + servers: [ + { + capabilities: 0, + created_at: '2026-01-01T00:00:00Z', + features: [], + geo_manual: false, + has_token: true, + hidden: false, + id: 'srv-1', + name: 'srv-1', + protocol_version: 2, + updated_at: '2026-01-01T00:00:00Z', + weight: 0 + } + ] + }) + projectServerCatalog(qc, { kind: 'tags_changed', serverId: 'srv-1', tags: ['old'] }) const { result } = renderHook(() => useUpdateServerTags('srv-1'), { wrapper: Wrapper }) await result.current.mutateAsync(['b', 'a']) expect(qc.getQueryData(['server-tags', 'srv-1'])).toEqual(['a', 'b']) - expect((qc.getQueryData(['servers']) as Array<{ id: string; tags: string[] }>)[0].tags).toEqual(['a', 'b']) + expect(readLiveServers(qc)?.[0].tags).toEqual(['a', 'b']) }) }) diff --git a/apps/web/src/hooks/use-server-tags.ts b/apps/web/src/hooks/use-server-tags.ts index 3b64883b0..1806b3191 100644 --- a/apps/web/src/hooks/use-server-tags.ts +++ b/apps/web/src/hooks/use-server-tags.ts @@ -1,6 +1,6 @@ import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { api } from '@/lib/api-client' +import { projectServerCatalog } from '@/lib/server-catalog' export function useServerTags(serverId: string, enabled = true) { return useQuery({ @@ -17,9 +17,7 @@ export function useUpdateServerTags(serverId: string) { mutationFn: (tags) => api.put(`/api/servers/${serverId}/tags`, { tags }), onSuccess: (data) => { queryClient.setQueryData(['server-tags', serverId], data) - queryClient.setQueryData(['servers'], (prev) => - prev?.map((s) => (s.id === serverId ? { ...s, tags: data } : s)) - ) + projectServerCatalog(queryClient, { kind: 'tags_changed', serverId, tags: data }) } }) } diff --git a/apps/web/src/hooks/use-servers-ws.test.ts b/apps/web/src/hooks/use-servers-ws.test.ts index 537bafd00..665c8e7ae 100644 --- a/apps/web/src/hooks/use-servers-ws.test.ts +++ b/apps/web/src/hooks/use-servers-ws.test.ts @@ -1,181 +1,12 @@ +import { QueryClient } from '@tanstack/react-query' import { describe, expect, it } from 'vitest' import { useUpgradeJobsStore } from '@/stores/upgrade-jobs-store' -import type { ServerMetrics } from './use-servers-ws' -import { - applyServerEdit, - handleWsMessage, - mergeServerUpdate, - setServerCapabilities, - setServerOnlineStatus -} from './use-servers-ws' - -function makeServer(overrides: Partial = {}): ServerMetrics { - return { - id: 's1', - name: 'Test', - online: true, - last_active: 0, - cpu: 50, - mem_used: 8_000_000_000, - mem_total: 16_000_000_000, - swap_used: 0, - swap_total: 4_000_000_000, - disk_used: 100_000_000_000, - disk_total: 500_000_000_000, - disk_read_bytes_per_sec: 0, - disk_write_bytes_per_sec: 0, - net_in_speed: 1000, - net_out_speed: 500, - net_in_transfer: 10_000, - net_out_transfer: 5000, - load1: 1.5, - load5: 1.2, - load15: 1.0, - tcp_conn: 100, - udp_conn: 10, - process_count: 200, - uptime: 3600, - cpu_name: 'Intel i7', - os: 'Linux', - region: 'US-East', - country_code: 'US', - group_id: 'g1', - ...overrides - } -} - -function makeQueryClient() { - const cache = new Map() - return { - setQueryData: (key: unknown[], value: unknown | ((prev: unknown) => unknown)) => { - const cacheKey = JSON.stringify(key) - const prev = cache.get(cacheKey) - const next = typeof value === 'function' ? (value as (prev: unknown) => unknown)(prev) : value - cache.set(cacheKey, next) - } - } -} - -describe('mergeServerUpdate', () => { - it('updates dynamic fields', () => { - const prev = [makeServer({ cpu: 50 })] - const incoming = [makeServer({ cpu: 75, mem_used: 10_000_000_000 })] - const result = mergeServerUpdate(prev, incoming) - expect(result[0].cpu).toBe(75) - expect(result[0].mem_used).toBe(10_000_000_000) - }) - - it('preserves static fields when incoming is null', () => { - const prev = [makeServer({ mem_total: 16_000_000_000, os: 'Linux', cpu_name: 'Intel i7' })] - const incoming = [makeServer({ id: 's1', mem_total: 0, os: null, cpu_name: null })] - const result = mergeServerUpdate(prev, incoming) - expect(result[0].mem_total).toBe(16_000_000_000) - expect(result[0].os).toBe('Linux') - expect(result[0].cpu_name).toBe('Intel i7') - }) - - it('preserves static fields when incoming is 0', () => { - const prev = [makeServer({ disk_total: 500_000_000_000, swap_total: 4_000_000_000 })] - const incoming = [makeServer({ id: 's1', disk_total: 0, swap_total: 0 })] - const result = mergeServerUpdate(prev, incoming) - expect(result[0].disk_total).toBe(500_000_000_000) - expect(result[0].swap_total).toBe(4_000_000_000) - }) - - it('ignores updates for unknown server id', () => { - const prev = [makeServer({ id: 's1' })] - const incoming = [makeServer({ id: 'unknown', cpu: 99 })] - const result = mergeServerUpdate(prev, incoming) - expect(result).toEqual(prev) - }) - - it('returns copy of prev when incoming is empty', () => { - const prev = [makeServer()] - const result = mergeServerUpdate(prev, []) - expect(result).toEqual(prev) - }) -}) - -describe('setServerOnlineStatus', () => { - it('sets target server offline', () => { - const prev = [makeServer({ id: 's1', online: true }), makeServer({ id: 's2', online: true })] - const result = setServerOnlineStatus(prev, 's1', false) - expect(result[0].online).toBe(false) - expect(result[1].online).toBe(true) - }) - - it('sets target server online', () => { - const prev = [makeServer({ id: 's1', online: false })] - const result = setServerOnlineStatus(prev, 's1', true) - expect(result[0].online).toBe(true) - }) - - it('leaves all unchanged for unknown server', () => { - const prev = [makeServer({ id: 's1', online: true })] - const result = setServerOnlineStatus(prev, 'unknown', false) - expect(result[0].online).toBe(true) - }) -}) - -describe('applyServerEdit', () => { - it('patches name and group_id of the edited server in place', () => { - const prev = [ - makeServer({ id: 's1', name: 'Old', group_id: null }), - makeServer({ id: 's2', name: 'Keep', group_id: 'g0' }) - ] - const result = applyServerEdit(prev, 's1', { name: 'New', group_id: 'g1', country_code: 'US' }) - expect(result[0].name).toBe('New') - expect(result[0].group_id).toBe('g1') - expect(result[0].country_code).toBe('US') - expect(result[1].name).toBe('Keep') - expect(result[1].group_id).toBe('g0') - }) - - it('preserves live metrics and online state while editing config', () => { - const prev = [makeServer({ id: 's1', name: 'Old', online: true, cpu: 42 })] - const result = applyServerEdit(prev, 's1', { name: 'Renamed', group_id: null, country_code: null }) - expect(result[0].online).toBe(true) - expect(result[0].cpu).toBe(42) - expect(result[0].group_id).toBeNull() - }) - - it('returns the list unchanged for an unknown server', () => { - const prev = [makeServer({ id: 's1', name: 'Old' })] - const result = applyServerEdit(prev, 'missing', { name: 'X', group_id: null, country_code: null }) - expect(result).toEqual(prev) - }) -}) - -describe('setServerCapabilities', () => { - it('updates configured, local, and effective capabilities together', () => { - const prev = [makeServer({ id: 's1' })] - const result = setServerCapabilities(prev, 's1', 64, 0, 0) - - expect(result[0].capabilities).toBe(64) - expect(result[0].agent_local_capabilities).toBe(0) - expect(result[0].effective_capabilities).toBe(0) - }) - - it('defaults temporary grants to an empty array when omitted', () => { - const prev = [makeServer({ id: 's1' })] - const result = setServerCapabilities(prev, 's1', 64, 0, 0) - - expect(result[0].temporary).toEqual([]) - }) - - it('merges temporary grant metadata onto the matching server', () => { - const prev = [makeServer({ id: 's1' })] - const temporary = [{ cap: 'CAP_TERMINAL', granted_at: 1000, expires_at: 2000 }] - const result = setServerCapabilities(prev, 's1', 64, 0, 0, temporary) - - expect(result[0].temporary).toEqual(temporary) - }) -}) +import { handleWsMessage } from './use-servers-ws' describe('handleWsMessage upgrade messages', () => { it('hydrates upgrade jobs from full_sync', () => { useUpgradeJobsStore.setState({ jobs: new Map() }) - const queryClient = makeQueryClient() + const queryClient = new QueryClient() handleWsMessage( { @@ -195,7 +26,7 @@ describe('handleWsMessage upgrade messages', () => { } ] }, - queryClient as never + queryClient ) expect(useUpgradeJobsStore.getState().jobs.get('server-1')?.job_id).toBe('job-1') @@ -220,7 +51,7 @@ describe('handleWsMessage upgrade messages', () => { ] ]) }) - const queryClient = makeQueryClient() + const queryClient = new QueryClient() handleWsMessage( { @@ -230,7 +61,7 @@ describe('handleWsMessage upgrade messages', () => { target_version: '1.2.3', stage: 'installing' }, - queryClient as never + queryClient ) expect(useUpgradeJobsStore.getState().jobs.get('server-1')?.stage).toBe('installing') @@ -238,7 +69,7 @@ describe('handleWsMessage upgrade messages', () => { it('stores terminal upgrade result from upgrade_result', () => { useUpgradeJobsStore.setState({ jobs: new Map() }) - const queryClient = makeQueryClient() + const queryClient = new QueryClient() handleWsMessage( { @@ -251,7 +82,7 @@ describe('handleWsMessage upgrade messages', () => { error: 'install failed', backup_path: '/tmp/backup' }, - queryClient as never + queryClient ) const job = useUpgradeJobsStore.getState().jobs.get('server-1') @@ -261,73 +92,3 @@ describe('handleWsMessage upgrade messages', () => { expect(job?.finished_at).not.toBeNull() }) }) - -function baseServer(overrides: Partial = {}): ServerMetrics { - return { - id: 'srv-1', - name: 'srv', - online: true, - country_code: null, - cpu: 0, - cpu_name: null, - cpu_cores: null, - disk_read_bytes_per_sec: 0, - disk_total: 0, - disk_used: 0, - disk_write_bytes_per_sec: 0, - group_id: null, - last_active: 0, - load1: 0, - load5: 0, - load15: 0, - mem_total: 0, - mem_used: 0, - net_in_speed: 0, - net_in_transfer: 0, - net_out_speed: 0, - net_out_transfer: 0, - os: null, - process_count: 0, - region: null, - swap_total: 0, - swap_used: 0, - tags: [], - tcp_conn: 0, - udp_conn: 0, - uptime: 0, - features: [], - ...overrides - } -} - -describe('mergeServerUpdate static-fields guard', () => { - it('preserves prior tags when incoming frame carries tags: []', () => { - const prev = [baseServer({ tags: ['prod', 'web'] })] - const incoming = [baseServer({ tags: [], cpu: 42 })] - const result = mergeServerUpdate(prev, incoming) - expect(result[0].tags).toEqual(['prod', 'web']) - expect(result[0].cpu).toBe(42) - }) - - it('preserves prior features when incoming frame carries features: []', () => { - const prev = [baseServer({ features: ['docker'] })] - const incoming = [baseServer({ features: [], cpu: 10 })] - const result = mergeServerUpdate(prev, incoming) - expect(result[0].features).toEqual(['docker']) - expect(result[0].cpu).toBe(10) - }) - - it('preserves prior cpu_cores when incoming frame carries cpu_cores: null', () => { - const prev = [baseServer({ cpu_cores: 8 })] - const incoming = [baseServer({ cpu_cores: null, cpu: 5 })] - const result = mergeServerUpdate(prev, incoming) - expect(result[0].cpu_cores).toBe(8) - }) - - it('overwrites prior tags with non-empty incoming array', () => { - const prev = [baseServer({ tags: ['old'] })] - const incoming = [baseServer({ tags: ['new-a', 'new-b'] })] - const result = mergeServerUpdate(prev, incoming) - expect(result[0].tags).toEqual(['new-a', 'new-b']) - }) -}) diff --git a/apps/web/src/hooks/use-servers-ws.ts b/apps/web/src/hooks/use-servers-ws.ts index 682b7a161..5dfdc4bdc 100644 --- a/apps/web/src/hooks/use-servers-ws.ts +++ b/apps/web/src/hooks/use-servers-ws.ts @@ -3,9 +3,10 @@ import { useQueryClient } from '@tanstack/react-query' import i18next from 'i18next' import { useEffect, useRef } from 'react' import { toast } from 'sonner' -import type { OutstandingEnrollmentSummary, SecurityEventDto, SecurityEventList } from '@/lib/api-schema' +import type { SecurityEventDto, SecurityEventList } from '@/lib/api-schema' import type { IpQualitySnapshotData, ServerIpQualityData, UnlockResultDto, UnlockStatus } from '@/lib/ip-quality-types' import type { NetworkProbeResultData } from '@/lib/network-types' +import { projectServerCatalog, type ServerMetrics } from '@/lib/server-catalog' import { WsClient } from '@/lib/ws-client' import type { DockerContainer, @@ -17,56 +18,6 @@ import { type UpgradeJob, useUpgradeJobsStore } from '@/stores/upgrade-jobs-stor const MAX_DOCKER_EVENTS = 100 const MAX_SECURITY_EVENTS_IN_CACHE = 200 -interface ServerMetrics { - agent_local_capabilities?: number | null - capabilities?: number - country_code: string | null - cpu: number - cpu_cores?: number | null - cpu_name: string | null - disk_read_bytes_per_sec: number - disk_total: number - disk_used: number - disk_write_bytes_per_sec: number - effective_capabilities?: number | null - features?: string[] - group_id: string | null - /** - * `true` iff the server row has a non-NULL `token_hash`. Pending servers (created via - * `POST /api/servers` but not yet enrolled by an agent) have `has_token = false`. - */ - has_token?: boolean - id: string - last_active: number - load1: number - load5: number - load15: number - mem_total: number - mem_used: number - name: string - net_in_speed: number - net_in_transfer: number - net_out_speed: number - net_out_transfer: number - online: boolean - os: string | null - /** - * Summary of the active outstanding enrollment (if any). Present for pending servers - * so the UI can render the install hint without an extra fetch. - */ - outstanding_enrollment?: OutstandingEnrollmentSummary | null - process_count: number - protocol_version?: number - region: string | null - swap_total: number - swap_used: number - tags?: string[] - tcp_conn: number - temporary?: Array<{ cap: string; granted_at: number; expires_at: number }> - udp_conn: number - uptime: number -} - type WsMessage = | { type: 'full_sync'; servers: ServerMetrics[]; upgrades?: UpgradeJob[] } | { type: 'update'; servers: ServerMetrics[] } @@ -137,193 +88,6 @@ interface WsUnlockResult { status: UnlockStatus } -export type { ServerMetrics } - -const STATIC_FIELDS = new Set([ - 'mem_total', - 'swap_total', - 'disk_total', - 'cpu_name', - 'cpu_cores', - 'os', - 'region', - 'country_code', - 'group_id', - 'tags', - 'features' -]) - -export function mergeServerUpdate(prev: ServerMetrics[], incoming: ServerMetrics[]): ServerMetrics[] { - const updated = [...prev] - // Index by id once (O(n)) instead of findIndex per incoming row: this runs on - // every WS `update`, which carries metrics for all servers, so the naive scan - // was O(n*m) each tick. The loop only updates rows in place, so indices stay valid. - const indexById = new Map(updated.map((s, i) => [s.id, i])) - for (const server of incoming) { - const idx = indexById.get(server.id) - if (idx !== undefined) { - const merged = { ...updated[idx] } - for (const [key, value] of Object.entries(server)) { - const isStaticDefault = - STATIC_FIELDS.has(key) && (value === null || value === 0 || (Array.isArray(value) && value.length === 0)) - if (!isStaticDefault) { - ;(merged as Record)[key] = value - } - } - updated[idx] = merged as ServerMetrics - } - } - return updated -} - -/** - * Default-zero runtime metrics for a brand-new server stub. The REST - * `ServerResponse` only carries static + identity fields; runtime metrics - * (cpu, mem_used, online, …) populate on the next WS push. - */ -function blankServerMetrics(id: string): ServerMetrics { - return { - id, - name: '', - cpu: 0, - cpu_cores: null, - cpu_name: null, - mem_total: 0, - mem_used: 0, - swap_total: 0, - swap_used: 0, - disk_total: 0, - disk_used: 0, - disk_read_bytes_per_sec: 0, - disk_write_bytes_per_sec: 0, - net_in_speed: 0, - net_in_transfer: 0, - net_out_speed: 0, - net_out_transfer: 0, - load1: 0, - load5: 0, - load15: 0, - process_count: 0, - tcp_conn: 0, - udp_conn: 0, - uptime: 0, - last_active: 0, - online: false, - country_code: null, - group_id: null, - os: null, - region: null - } -} - -/** - * Reconcile the WS-fed `['servers']` cache against a fresh REST list. Membership - * comes from REST (handles add / delete / cleanup), but per-row runtime metrics - * are preserved from the existing cache so we don't blink to zeros while - * waiting for the next WS update. - */ -export function reconcileServersFromRest( - prev: ServerMetrics[] | undefined, - fresh: Array & { id: string }> -): ServerMetrics[] { - const prevMap = new Map((prev ?? []).map((s) => [s.id, s])) - return fresh.map((row) => { - const existing = prevMap.get(row.id) - if (existing) { - const merged: Record = { ...existing } - for (const [key, value] of Object.entries(row)) { - const isStaticDefault = - STATIC_FIELDS.has(key) && (value === null || value === 0 || (Array.isArray(value) && value.length === 0)) - if (!isStaticDefault) { - merged[key] = value - } - } - return merged as unknown as ServerMetrics - } - return { ...blankServerMetrics(row.id), ...row } as ServerMetrics - }) -} - -export function setServerOnlineStatus(prev: ServerMetrics[], serverId: string, online: boolean): ServerMetrics[] { - return prev.map((s) => (s.id === serverId ? { ...s, online } : s)) -} - -export function setServerDockerAvailability( - prev: ServerMetrics[], - serverId: string, - available: boolean -): ServerMetrics[] { - return prev.map((s) => { - if (s.id !== serverId) { - return s - } - const features = s.features ?? [] - if (available && !features.includes('docker')) { - return { ...s, features: [...features, 'docker'] } - } - if (!available && features.includes('docker')) { - return { ...s, features: features.filter((f) => f !== 'docker') } - } - return s - }) -} - -export function applyServerEdit( - prev: ServerMetrics[], - serverId: string, - edited: { name: string; group_id: string | null; country_code: string | null } -): ServerMetrics[] { - return prev.map((server) => - server.id === serverId - ? { ...server, name: edited.name, group_id: edited.group_id, country_code: edited.country_code } - : server - ) -} - -export function setServerCapabilities( - prev: ServerMetrics[], - serverId: string, - capabilities: number, - agentLocalCapabilities: number | null | undefined, - effectiveCapabilities: number | null | undefined, - temporary?: Array<{ cap: string; granted_at: number; expires_at: number }> | null -): ServerMetrics[] { - return prev.map((server) => - server.id === serverId - ? { - ...server, - capabilities, - agent_local_capabilities: agentLocalCapabilities ?? null, - effective_capabilities: effectiveCapabilities ?? null, - temporary: temporary ?? [] - } - : server - ) -} - -function setServerDetailDockerAvailability( - prev: Record | undefined, - available: boolean -): Record | undefined { - if (!prev) { - return prev - } - - const features = Array.isArray(prev.features) - ? prev.features.filter((feature): feature is string => typeof feature === 'string') - : [] - - if (available && !features.includes('docker')) { - return { ...prev, features: [...features, 'docker'] } - } - - if (!available && features.includes('docker')) { - return { ...prev, features: features.filter((feature) => feature !== 'docker') } - } - - return prev -} - type QueryClient = ReturnType type FullSyncMessage = Extract type UpdateMessage = Extract @@ -333,16 +97,14 @@ function isWsMessageLike(raw: unknown): raw is { type: string } & Record(['servers'], msg.servers) + projectServerCatalog(queryClient, { kind: 'ws_full_sync', servers: msg.servers }) if (Array.isArray(msg.upgrades)) { - useUpgradeJobsStore.getState().setJobs(msg.upgrades as UpgradeJob[]) + useUpgradeJobsStore.getState().setJobs(msg.upgrades) } } function handleUpdateMessage(msg: UpdateMessage, queryClient: QueryClient): void { - queryClient.setQueryData(['servers'], (prev) => - prev ? mergeServerUpdate(prev, msg.servers) : msg.servers - ) + projectServerCatalog(queryClient, { kind: 'ws_update', servers: msg.servers }) } function handleServerMetricsMessage(raw: { type: string } & Record, queryClient: QueryClient): void { @@ -364,9 +126,7 @@ function handleServerMetricsMessage(raw: { type: string } & Record(['servers'], (prev) => - prev ? setServerOnlineStatus(prev, server_id, online) : prev - ) + projectServerCatalog(queryClient, { kind: 'online_changed', serverId: server_id, online }) } } @@ -377,42 +137,14 @@ function handleCapabilityMessage(raw: { type: string } & Record } const msg = raw as WsMessage & { type: 'capabilities_changed' } const { server_id, capabilities, agent_local_capabilities, effective_capabilities, temporary } = msg - queryClient.setQueryData(['servers'], (prev) => - prev - ? setServerCapabilities( - prev, - server_id, - capabilities, - agent_local_capabilities, - effective_capabilities, - temporary - ) - : prev - ) - queryClient.setQueryData(['servers', server_id], (prev: Record | undefined) => - prev - ? { - ...prev, - capabilities, - agent_local_capabilities: agent_local_capabilities ?? null, - effective_capabilities: effective_capabilities ?? null, - temporary: temporary ?? [] - } - : prev - ) - queryClient.setQueryData[]>(['servers-list'], (prev) => - prev?.map((s) => - s.id === server_id - ? { - ...s, - capabilities, - agent_local_capabilities: agent_local_capabilities ?? null, - effective_capabilities: effective_capabilities ?? null, - temporary: temporary ?? [] - } - : s - ) - ) + projectServerCatalog(queryClient, { + kind: 'capabilities_changed', + serverId: server_id, + capabilities, + agentLocalCapabilities: agent_local_capabilities, + effectiveCapabilities: effective_capabilities, + temporary + }) return } if (raw.type === 'agent_info_updated') { @@ -421,12 +153,12 @@ function handleCapabilityMessage(raw: { type: string } & Record } const msg = raw as WsMessage & { type: 'agent_info_updated' } const { server_id, protocol_version, agent_version } = msg - queryClient.setQueryData(['servers', server_id], (prev: Record | undefined) => - prev ? { ...prev, protocol_version, agent_version: agent_version ?? null } : prev - ) - queryClient.setQueryData[]>(['servers-list'], (prev) => - prev?.map((s) => (s.id === server_id ? { ...s, protocol_version, agent_version: agent_version ?? null } : s)) - ) + projectServerCatalog(queryClient, { + kind: 'agent_info_changed', + serverId: server_id, + protocolVersion: protocol_version, + agentVersion: agent_version + }) } } @@ -466,12 +198,7 @@ function handleDockerMessage(raw: { type: string } & Record, qu } const msg = raw as WsMessage & { type: 'docker_availability_changed' } const { server_id, available } = msg - queryClient.setQueryData(['servers'], (prev) => - prev ? setServerDockerAvailability(prev, server_id, available) : prev - ) - queryClient.setQueryData(['servers', server_id], (prev: Record | undefined) => - setServerDetailDockerAvailability(prev, available) - ) + projectServerCatalog(queryClient, { kind: 'docker_availability_changed', serverId: server_id, available }) } } diff --git a/apps/web/src/hooks/use-upgrade-job.ts b/apps/web/src/hooks/use-upgrade-job.ts index 8a3386ce7..18ed1c29b 100644 --- a/apps/web/src/hooks/use-upgrade-job.ts +++ b/apps/web/src/hooks/use-upgrade-job.ts @@ -1,6 +1,7 @@ import { useMutation, useQueryClient } from '@tanstack/react-query' import { api } from '@/lib/api-client' import type { UpgradeRequest } from '@/lib/api-schema' +import { invalidateServerDetail } from '@/lib/server-catalog' import { type UpgradeJob, useUpgradeJobsStore } from '@/stores/upgrade-jobs-store' interface UseUpgradeJobResult { @@ -19,7 +20,7 @@ export function useUpgradeJob(serverId: string): UseUpgradeJobResult { return api.post(`/api/servers/${serverId}/upgrade`, body) }, onSuccess: () => { - queryClient.invalidateQueries({ queryKey: ['servers', serverId] }) + invalidateServerDetail(queryClient, serverId).catch(() => undefined) } }) diff --git a/apps/web/src/lib/dev-mock-servers.ts b/apps/web/src/lib/dev-mock-servers.ts index 0b659ffa7..279a3e048 100644 --- a/apps/web/src/lib/dev-mock-servers.ts +++ b/apps/web/src/lib/dev-mock-servers.ts @@ -1,4 +1,4 @@ -import type { ServerMetrics } from '@/hooks/use-servers-ws' +import type { ServerMetrics } from '@/lib/server-catalog' const STORAGE_KEY = 'serverbee-mock-servers' const GIB = 1024 ** 3 diff --git a/apps/web/src/lib/server-catalog.test.ts b/apps/web/src/lib/server-catalog.test.ts new file mode 100644 index 000000000..ebdeb3065 --- /dev/null +++ b/apps/web/src/lib/server-catalog.test.ts @@ -0,0 +1,552 @@ +import { QueryClient, QueryClientProvider } from '@tanstack/react-query' +import { renderHook, waitFor } from '@testing-library/react' +import { createElement, type ReactNode } from 'react' +import { describe, expect, it, vi } from 'vitest' +import type { OutstandingEnrollmentSummary, ServerResponse } from '@/lib/api-schema' +import { + invalidateServerDetail, + projectServerCatalog, + readLiveServers, + refreshServerCatalog, + type ServerMetrics, + subscribeLiveServers, + useLiveServers, + useServerDetail, + useServerList +} from './server-catalog' + +function createQueryClient(): QueryClient { + return new QueryClient({ + defaultOptions: { + queries: { retry: false } + } + }) +} + +function createWrapper(queryClient: QueryClient) { + return function QueryWrapper({ children }: { children: ReactNode }) { + return createElement(QueryClientProvider, { client: queryClient }, children) + } +} + +function apiResponse(data: T): Response { + return new Response(JSON.stringify({ data }), { + headers: { 'Content-Type': 'application/json' }, + status: 200 + }) +} + +function deferred(): { promise: Promise; resolve: (value: T) => void } { + let resolvePromise: ((value: T) => void) | null = null + const promise = new Promise((resolve) => { + resolvePromise = resolve + }) + return { + promise, + resolve: (value) => { + if (!resolvePromise) { + throw new Error('Deferred promise was not initialized') + } + resolvePromise(value) + } + } +} + +function readServerList(queryClient: QueryClient): ServerResponse[] | undefined { + const { result, unmount } = renderHook(() => useServerList({ enabled: false }), { + wrapper: createWrapper(queryClient) + }) + const data = result.current.data + unmount() + return data +} + +function readServerDetail(queryClient: QueryClient, serverId: string): ServerResponse | undefined { + const { result, unmount } = renderHook(() => useServerDetail(serverId, { enabled: false }), { + wrapper: createWrapper(queryClient) + }) + const data = result.current.data + unmount() + return data +} + +function makeRestServer(overrides: Partial = {}): ServerResponse { + return { + capabilities: 1852, + country_code: 'US', + cpu_cores: 8, + cpu_name: 'AMD EPYC', + created_at: '2026-07-01T00:00:00Z', + disk_total: 500_000, + effective_capabilities: 1852, + features: ['docker'], + geo_manual: false, + group_id: 'edge', + has_token: true, + hidden: false, + id: 'server-1', + mem_total: 32_000, + name: 'Edge One', + os: 'Linux', + outstanding_enrollment: null, + protocol_version: 2, + region: 'us-east', + swap_total: 4000, + temporary: [], + updated_at: '2026-07-01T00:00:00Z', + weight: 100, + ...overrides + } +} + +function makeMetrics(overrides: Partial = {}): ServerMetrics { + return { + country_code: 'US', + cpu: 42, + cpu_cores: 8, + cpu_name: 'AMD EPYC', + disk_read_bytes_per_sec: 100, + disk_total: 500_000, + disk_used: 125_000, + disk_write_bytes_per_sec: 50, + features: ['docker'], + group_id: 'edge', + has_token: true, + id: 'server-1', + last_active: 1000, + load1: 1, + load5: 0.8, + load15: 0.6, + mem_total: 32_000, + mem_used: 16_000, + name: 'Edge One', + net_in_speed: 1000, + net_in_transfer: 10_000, + net_out_speed: 500, + net_out_transfer: 5000, + online: true, + os: 'Linux', + process_count: 120, + region: 'us-east', + swap_total: 4000, + swap_used: 100, + tags: ['prod'], + tcp_conn: 30, + udp_conn: 5, + uptime: 3600, + ...overrides + } +} + +const OUTSTANDING_ENROLLMENT: OutstandingEnrollmentSummary = { + code_prefix: 'abcdef', + created_at: '2026-07-11T00:00:00Z', + expires_at: '2026-07-11T00:10:00Z', + id: 'enrollment-1' +} + +describe('server catalog projection', () => { + it('preserves WS runtime data while an authoritative REST snapshot applies null clears and membership', () => { + const queryClient = createQueryClient() + const live = makeMetrics({ cpu: 73, mem_used: 20_000, online: true }) + const priorDetail = makeRestServer() + const cleared = makeRestServer({ + country_code: null, + cpu_cores: null, + cpu_name: null, + disk_total: null, + features: [], + group_id: null, + mem_total: null, + name: 'Renamed Edge', + os: null, + region: null, + swap_total: null + }) + const added = makeRestServer({ id: 'server-2', name: 'Pending Two' }) + + projectServerCatalog(queryClient, { kind: 'ws_full_sync', servers: [live] }) + projectServerCatalog(queryClient, { kind: 'server_saved', server: priorDetail }) + projectServerCatalog(queryClient, { kind: 'rest_snapshot', servers: [cleared, added] }) + + const projected = readLiveServers(queryClient) + expect(projected).toHaveLength(2) + expect(projected?.[0]).toMatchObject({ + country_code: null, + cpu: 73, + cpu_cores: null, + cpu_name: null, + disk_total: 0, + features: [], + group_id: null, + mem_total: 0, + mem_used: 20_000, + name: 'Renamed Edge', + online: true, + os: null, + region: null, + swap_total: 0 + }) + expect(projected?.[1]).toMatchObject({ cpu: 0, id: 'server-2', name: 'Pending Two', online: false }) + expect(readServerList(queryClient)).toEqual([cleared, added]) + expect(readServerDetail(queryClient, 'server-1')).toEqual(cleared) + }) + + it('updates runtime metrics without letting WS placeholders erase REST or mutation-owned fields', () => { + const queryClient = createQueryClient() + const rest = makeRestServer({ + has_token: false, + name: 'Edited Name', + outstanding_enrollment: OUTSTANDING_ENROLLMENT + }) + + projectServerCatalog(queryClient, { kind: 'rest_snapshot', servers: [rest] }) + projectServerCatalog(queryClient, { kind: 'tags_changed', serverId: rest.id, tags: ['prod', 'edge'] }) + projectServerCatalog(queryClient, { + kind: 'ws_update', + servers: [ + makeMetrics({ + country_code: null, + cpu: 91, + cpu_cores: null, + cpu_name: null, + disk_total: 0, + features: [], + group_id: null, + has_token: true, + mem_total: 0, + mem_used: 24_000, + name: 'Stale Agent Name', + os: null, + outstanding_enrollment: null, + region: null, + swap_total: 0, + tags: [] + }) + ] + }) + + expect(readLiveServers(queryClient)?.[0]).toMatchObject({ + country_code: 'US', + cpu: 91, + cpu_cores: 8, + cpu_name: 'AMD EPYC', + disk_total: 500_000, + features: ['docker'], + group_id: 'edge', + has_token: false, + mem_total: 32_000, + mem_used: 24_000, + name: 'Edited Name', + os: 'Linux', + outstanding_enrollment: OUTSTANDING_ENROLLMENT, + region: 'us-east', + swap_total: 4000, + tags: ['prod', 'edge'] + }) + }) + + it('uses full sync as authoritative membership and static state while preserving out-of-band metadata', () => { + const queryClient = createQueryClient() + const first = makeRestServer({ + disk_total: null, + mem_total: null, + outstanding_enrollment: OUTSTANDING_ENROLLMENT, + swap_total: null + }) + const second = makeRestServer({ id: 'server-2', name: 'REST Only' }) + + projectServerCatalog(queryClient, { kind: 'rest_snapshot', servers: [first, second] }) + projectServerCatalog(queryClient, { kind: 'server_saved', server: first }) + projectServerCatalog(queryClient, { kind: 'server_saved', server: second }) + projectServerCatalog(queryClient, { + agentLocalCapabilities: 64, + capabilities: 64, + effectiveCapabilities: 64, + kind: 'capabilities_changed', + serverId: first.id, + temporary: [] + }) + projectServerCatalog(queryClient, { + kind: 'ws_full_sync', + servers: [ + makeMetrics({ + country_code: null, + cpu: 87, + cpu_cores: null, + cpu_name: null, + disk_total: 0, + features: [], + group_id: null, + has_token: false, + mem_total: 0, + name: 'Authoritative Full Sync', + os: null, + outstanding_enrollment: null, + region: null, + swap_total: 0, + tags: [] + }) + ] + }) + + const projected = readLiveServers(queryClient) + expect(projected?.map((server) => server.id)).toEqual([first.id]) + expect(projected?.[0]).toMatchObject({ + capabilities: 64, + country_code: null, + cpu: 87, + cpu_cores: null, + cpu_name: null, + disk_total: 0, + features: [], + group_id: null, + has_token: false, + mem_total: 0, + name: 'Authoritative Full Sync', + os: null, + outstanding_enrollment: null, + region: null, + swap_total: 0, + tags: [] + }) + expect(readServerList(queryClient)?.map((server) => server.id)).toEqual([first.id]) + expect(readServerList(queryClient)?.[0]).toMatchObject({ + country_code: null, + disk_total: null, + features: [], + has_token: false, + mem_total: null, + name: 'Authoritative Full Sync', + swap_total: null + }) + expect(readServerDetail(queryClient, second.id)).toBeUndefined() + }) + + it('propagates capability, agent-info, and Docker changes to live, list, and existing detail projections', () => { + const queryClient = createQueryClient() + const rest = makeRestServer({ features: [] }) + const temporary = [{ cap: 'CAP_TERMINAL', expires_at: 2000, granted_at: 1000 }] + + projectServerCatalog(queryClient, { kind: 'server_saved', server: rest }) + projectServerCatalog(queryClient, { kind: 'rest_snapshot', servers: [rest] }) + projectServerCatalog(queryClient, { + agentLocalCapabilities: 64, + capabilities: 64, + effectiveCapabilities: 65, + kind: 'capabilities_changed', + serverId: rest.id, + temporary + }) + projectServerCatalog(queryClient, { + agentVersion: '1.4.0', + kind: 'agent_info_changed', + protocolVersion: 3, + serverId: rest.id + }) + projectServerCatalog(queryClient, { + available: true, + kind: 'docker_availability_changed', + serverId: rest.id + }) + + const expected = { + agent_local_capabilities: 64, + agent_version: '1.4.0', + capabilities: 64, + effective_capabilities: 65, + features: ['docker'], + protocol_version: 3, + temporary + } + expect(readLiveServers(queryClient)?.[0]).toMatchObject(expected) + expect(readServerList(queryClient)?.[0]).toMatchObject(expected) + expect(readServerDetail(queryClient, rest.id)).toMatchObject(expected) + + projectServerCatalog(queryClient, { + available: false, + kind: 'docker_availability_changed', + serverId: rest.id + }) + expect(readLiveServers(queryClient)?.[0].features).toEqual([]) + expect(readServerList(queryClient)?.[0].features).toEqual([]) + expect(readServerDetail(queryClient, rest.id)?.features).toEqual([]) + }) + + it('keeps edit, enrollment, tags, and removal outcomes consistent without erasing runtime data', () => { + const queryClient = createQueryClient() + const first = makeRestServer() + const second = makeRestServer({ id: 'server-2', name: 'Keep Me' }) + const saved = makeRestServer({ group_id: null, name: 'Edited Name' }) + + projectServerCatalog(queryClient, { kind: 'rest_snapshot', servers: [first, second] }) + projectServerCatalog(queryClient, { kind: 'online_changed', online: true, serverId: first.id }) + projectServerCatalog(queryClient, { kind: 'server_saved', server: saved }) + projectServerCatalog(queryClient, { kind: 'tags_changed', serverId: first.id, tags: ['edited'] }) + projectServerCatalog(queryClient, { + kind: 'enrollment_changed', + outstandingEnrollment: OUTSTANDING_ENROLLMENT, + serverId: first.id, + tokenRevoked: true + }) + + expect(readLiveServers(queryClient)?.[0]).toMatchObject({ + group_id: null, + has_token: false, + name: 'Edited Name', + online: false, + outstanding_enrollment: OUTSTANDING_ENROLLMENT, + tags: ['edited'] + }) + expect(readServerList(queryClient)?.[0]).toMatchObject({ + group_id: null, + has_token: false, + name: 'Edited Name', + outstanding_enrollment: OUTSTANDING_ENROLLMENT + }) + expect(readServerDetail(queryClient, first.id)).toMatchObject({ + has_token: false, + name: 'Edited Name', + outstanding_enrollment: OUTSTANDING_ENROLLMENT + }) + + projectServerCatalog(queryClient, { kind: 'servers_removed', serverIds: [first.id] }) + expect(readLiveServers(queryClient)?.map((server) => server.id)).toEqual([second.id]) + expect(readServerList(queryClient)?.map((server) => server.id)).toEqual([second.id]) + expect(readServerDetail(queryClient, first.id)).toBeUndefined() + }) + + it('removes detail projections missing from an authoritative REST snapshot', () => { + const queryClient = createQueryClient() + const removed = makeRestServer() + const kept = makeRestServer({ id: 'server-2', name: 'Kept' }) + + projectServerCatalog(queryClient, { kind: 'server_saved', server: removed }) + expect(readServerDetail(queryClient, removed.id)).toEqual(removed) + + projectServerCatalog(queryClient, { kind: 'rest_snapshot', servers: [kept] }) + expect(readServerDetail(queryClient, removed.id)).toBeUndefined() + }) + + it('notifies live subscribers only for the exact live projection', async () => { + const queryClient = createQueryClient() + const server = makeRestServer() + projectServerCatalog(queryClient, { kind: 'rest_snapshot', servers: [server] }) + projectServerCatalog(queryClient, { kind: 'server_saved', server }) + const listener = vi.fn() + const unsubscribe = subscribeLiveServers(queryClient, listener) + + const liveObserver = renderHook(() => useLiveServers({ enabled: false }), { + wrapper: createWrapper(queryClient) + }) + liveObserver.unmount() + + await invalidateServerDetail(queryClient, server.id) + queryClient.setQueryData(['servers', server.id, 'records', 24], [{ cpu: 1 }]) + expect(listener).not.toHaveBeenCalled() + + projectServerCatalog(queryClient, { kind: 'online_changed', online: true, serverId: server.id }) + expect(listener).toHaveBeenCalledTimes(1) + expect(listener).toHaveBeenLastCalledWith([expect.objectContaining({ id: server.id, online: true })]) + + unsubscribe() + }) + + it('retains unobserved live membership beyond the default query garbage-collection window', () => { + vi.useFakeTimers() + try { + const queryClient = createQueryClient() + const server = makeMetrics() + projectServerCatalog(queryClient, { kind: 'ws_full_sync', servers: [server] }) + + vi.advanceTimersByTime(6 * 60 * 1000) + + expect(readLiveServers(queryClient)?.map((entry) => entry.id)).toEqual([server.id]) + } finally { + vi.useRealTimers() + } + }) + + it('projects a normal REST list query through the catalog owner', async () => { + const queryClient = createQueryClient() + const server = makeRestServer({ id: 'server-from-http', name: 'Fetched Server' }) + const fetchMock = vi.spyOn(globalThis, 'fetch').mockResolvedValue(apiResponse([server])) + + const hook = renderHook(() => useServerList(), { wrapper: createWrapper(queryClient) }) + await waitFor(() => expect(hook.result.current.isSuccess).toBe(true)) + + expect(fetchMock).toHaveBeenCalledWith('/api/servers', expect.objectContaining({ method: 'GET' })) + expect(readLiveServers(queryClient)?.[0]).toMatchObject({ + cpu: 0, + id: server.id, + name: server.name, + online: false + }) + expect(hook.result.current.data).toEqual([server]) + + hook.unmount() + fetchMock.mockRestore() + }) + + it('projects a REST detail query through live and list views', async () => { + const queryClient = createQueryClient() + const stale = makeRestServer({ name: 'Stale Name' }) + const fresh = makeRestServer({ name: 'Fresh Name' }) + projectServerCatalog(queryClient, { kind: 'rest_snapshot', servers: [stale] }) + const fetchMock = vi.spyOn(globalThis, 'fetch').mockResolvedValue(apiResponse(fresh)) + + const hook = renderHook(() => useServerDetail(fresh.id), { wrapper: createWrapper(queryClient) }) + await waitFor(() => expect(hook.result.current.isSuccess).toBe(true)) + + expect(readLiveServers(queryClient)?.[0].name).toBe('Fresh Name') + expect(readServerList(queryClient)?.[0].name).toBe('Fresh Name') + expect(readServerDetail(queryClient, fresh.id)?.name).toBe('Fresh Name') + + hook.unmount() + fetchMock.mockRestore() + }) + + it('prevents an older overlapping REST refresh from replacing a newer snapshot', async () => { + const queryClient = createQueryClient() + const oldResponse = deferred() + const newResponse = deferred() + const fetchMock = vi + .spyOn(globalThis, 'fetch') + .mockImplementationOnce(() => oldResponse.promise) + .mockImplementationOnce(() => newResponse.promise) + + const oldRefresh = refreshServerCatalog(queryClient) + await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(1)) + const newRefresh = refreshServerCatalog(queryClient) + await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)) + + newResponse.resolve(apiResponse([makeRestServer({ name: 'New Snapshot' })])) + await newRefresh + oldResponse.resolve(apiResponse([makeRestServer({ name: 'Old Snapshot' })])) + await oldRefresh + + expect(readLiveServers(queryClient)?.[0].name).toBe('New Snapshot') + expect(readServerList(queryClient)?.[0].name).toBe('New Snapshot') + fetchMock.mockRestore() + }) + + it('forces an explicit refresh even while the cached list is globally fresh', async () => { + const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false, staleTime: 30_000 } } + }) + projectServerCatalog(queryClient, { + kind: 'rest_snapshot', + servers: [makeRestServer({ name: 'Cached Snapshot' })] + }) + const fetchMock = vi + .spyOn(globalThis, 'fetch') + .mockResolvedValue(apiResponse([makeRestServer({ name: 'Refreshed Snapshot' })])) + + await refreshServerCatalog(queryClient) + + expect(fetchMock).toHaveBeenCalledTimes(1) + expect(readLiveServers(queryClient)?.[0].name).toBe('Refreshed Snapshot') + expect(readServerList(queryClient)?.[0].name).toBe('Refreshed Snapshot') + fetchMock.mockRestore() + }) +}) diff --git a/apps/web/src/lib/server-catalog.ts b/apps/web/src/lib/server-catalog.ts new file mode 100644 index 000000000..f43cec607 --- /dev/null +++ b/apps/web/src/lib/server-catalog.ts @@ -0,0 +1,620 @@ +import { type QueryClient, useQuery, useQueryClient } from '@tanstack/react-query' +import { api } from '@/lib/api-client' +import type { OutstandingEnrollmentSummary, ServerResponse } from '@/lib/api-schema' + +const LIVE_SERVERS_KEY = ['server-catalog', 'live'] as const +const SERVER_LIST_KEY = ['server-catalog', 'list'] as const +const SERVER_DETAIL_PREFIX = ['server-catalog', 'detail'] as const +const listRequestGenerations = new WeakMap() + +type TemporaryGrant = NonNullable[number] + +interface CatalogQueryOptions { + enabled?: boolean +} + +export interface ServerMetrics { + agent_local_capabilities?: number | null + agent_version?: string | null + capabilities?: number + country_code: string | null + cpu: number + cpu_cores?: number | null + cpu_name: string | null + disk_read_bytes_per_sec: number + disk_total: number + disk_used: number + disk_write_bytes_per_sec: number + effective_capabilities?: number | null + features?: string[] + group_id: string | null + has_token?: boolean + id: string + last_active: number + load1: number + load5: number + load15: number + mem_total: number + mem_used: number + name: string + net_in_speed: number + net_in_transfer: number + net_out_speed: number + net_out_transfer: number + online: boolean + os: string | null + outstanding_enrollment?: OutstandingEnrollmentSummary | null + process_count: number + protocol_version?: number + region: string | null + swap_total: number + swap_used: number + tags?: string[] + tcp_conn: number + temporary?: TemporaryGrant[] + udp_conn: number + uptime: number +} + +export type ServerCatalogEvent = + | { kind: 'rest_snapshot'; servers: ServerResponse[] } + | { kind: 'servers_removed'; serverIds: readonly string[] } + | { kind: 'server_saved'; server: ServerResponse } + | { + kind: 'enrollment_changed' + outstandingEnrollment: OutstandingEnrollmentSummary | null + serverId: string + tokenRevoked: boolean + } + | { kind: 'tags_changed'; serverId: string; tags: string[] } + | { kind: 'ws_full_sync'; servers: ServerMetrics[] } + | { kind: 'ws_update'; servers: ServerMetrics[] } + | { kind: 'online_changed'; online: boolean; serverId: string } + | { + agentLocalCapabilities: number | null | undefined + capabilities: number + effectiveCapabilities: number | null | undefined + kind: 'capabilities_changed' + serverId: string + temporary: TemporaryGrant[] | undefined + } + | { + agentVersion: string | null | undefined + kind: 'agent_info_changed' + protocolVersion: number + serverId: string + } + | { available: boolean; kind: 'docker_availability_changed'; serverId: string } + +function serverDetailKey(serverId: string): readonly ['server-catalog', 'detail', string] { + return [...SERVER_DETAIL_PREFIX, serverId] +} + +function blankServerMetrics(id: string): ServerMetrics { + return { + agent_local_capabilities: null, + agent_version: null, + country_code: null, + cpu: 0, + cpu_cores: null, + cpu_name: null, + disk_read_bytes_per_sec: 0, + disk_total: 0, + disk_used: 0, + disk_write_bytes_per_sec: 0, + effective_capabilities: null, + features: [], + group_id: null, + has_token: false, + id, + last_active: 0, + load1: 0, + load5: 0, + load15: 0, + mem_total: 0, + mem_used: 0, + name: '', + net_in_speed: 0, + net_in_transfer: 0, + net_out_speed: 0, + net_out_transfer: 0, + online: false, + os: null, + outstanding_enrollment: null, + process_count: 0, + region: null, + swap_total: 0, + swap_used: 0, + tags: [], + tcp_conn: 0, + temporary: [], + udp_conn: 0, + uptime: 0 + } +} + +function projectRestServer(current: ServerMetrics, server: ServerResponse): ServerMetrics { + return { + ...current, + agent_local_capabilities: + server.agent_local_capabilities === undefined + ? current.agent_local_capabilities + : server.agent_local_capabilities, + agent_version: server.agent_version === undefined ? current.agent_version : server.agent_version, + capabilities: server.capabilities, + country_code: server.country_code === undefined ? current.country_code : server.country_code, + cpu_cores: server.cpu_cores === undefined ? current.cpu_cores : server.cpu_cores, + cpu_name: server.cpu_name === undefined ? current.cpu_name : server.cpu_name, + disk_total: server.disk_total === undefined ? current.disk_total : (server.disk_total ?? 0), + effective_capabilities: + server.effective_capabilities === undefined ? current.effective_capabilities : server.effective_capabilities, + features: [...server.features], + group_id: server.group_id === undefined ? current.group_id : server.group_id, + has_token: server.has_token, + mem_total: server.mem_total === undefined ? current.mem_total : (server.mem_total ?? 0), + name: server.name, + os: server.os === undefined ? current.os : server.os, + outstanding_enrollment: + server.outstanding_enrollment === undefined ? current.outstanding_enrollment : server.outstanding_enrollment, + protocol_version: server.protocol_version, + region: server.region === undefined ? current.region : server.region, + swap_total: server.swap_total === undefined ? current.swap_total : (server.swap_total ?? 0), + temporary: server.temporary === undefined ? current.temporary : [...server.temporary] + } +} + +function projectRestSnapshot(current: ServerMetrics[] | undefined, servers: ServerResponse[]): ServerMetrics[] { + const currentById = new Map((current ?? []).map((server) => [server.id, server])) + return servers.map((server) => projectRestServer(currentById.get(server.id) ?? blankServerMetrics(server.id), server)) +} + +function projectFullSyncServerToRest(current: ServerResponse, server: ServerMetrics): ServerResponse { + return { + ...current, + country_code: server.country_code, + cpu_cores: server.cpu_cores, + cpu_name: server.cpu_name, + features: server.features === undefined ? current.features : [...server.features], + group_id: server.group_id, + has_token: server.has_token === undefined ? current.has_token : server.has_token, + name: server.name, + os: server.os, + outstanding_enrollment: + server.outstanding_enrollment === undefined ? current.outstanding_enrollment : server.outstanding_enrollment, + region: server.region + } +} + +function mergeWsServerUpdate(current: ServerMetrics, incoming: ServerMetrics): ServerMetrics { + return { + ...current, + cpu: incoming.cpu, + disk_read_bytes_per_sec: incoming.disk_read_bytes_per_sec, + disk_used: incoming.disk_used, + disk_write_bytes_per_sec: incoming.disk_write_bytes_per_sec, + last_active: incoming.last_active, + load1: incoming.load1, + load5: incoming.load5, + load15: incoming.load15, + mem_used: incoming.mem_used, + net_in_speed: incoming.net_in_speed, + net_in_transfer: incoming.net_in_transfer, + net_out_speed: incoming.net_out_speed, + net_out_transfer: incoming.net_out_transfer, + online: incoming.online, + process_count: incoming.process_count, + swap_used: incoming.swap_used, + tcp_conn: incoming.tcp_conn, + udp_conn: incoming.udp_conn, + uptime: incoming.uptime + } +} + +function mergeWsUpdate(current: ServerMetrics[], incoming: ServerMetrics[]): ServerMetrics[] { + const incomingById = new Map(incoming.map((server) => [server.id, server])) + return current.map((server) => { + const update = incomingById.get(server.id) + return update ? mergeWsServerUpdate(server, update) : server + }) +} + +function mergeWsFullSync(current: ServerMetrics[] | undefined, incoming: ServerMetrics[]): ServerMetrics[] { + const currentById = new Map((current ?? []).map((server) => [server.id, server])) + return incoming.map((server) => { + const existing = currentById.get(server.id) + if (!existing) { + return server + } + return { + ...server, + agent_local_capabilities: + server.agent_local_capabilities === undefined + ? existing.agent_local_capabilities + : server.agent_local_capabilities, + agent_version: server.agent_version === undefined ? existing.agent_version : server.agent_version, + capabilities: server.capabilities === undefined ? existing.capabilities : server.capabilities, + effective_capabilities: + server.effective_capabilities === undefined ? existing.effective_capabilities : server.effective_capabilities, + protocol_version: server.protocol_version === undefined ? existing.protocol_version : server.protocol_version, + temporary: server.temporary === undefined ? existing.temporary : server.temporary + } + }) +} + +function updateById( + rows: T[] | undefined, + serverId: string, + update: (row: T) => T +): T[] | undefined { + if (!rows) { + return rows + } + const index = rows.findIndex((row) => row.id === serverId) + if (index < 0) { + return rows + } + const updated = update(rows[index]) + if (updated === rows[index]) { + return rows + } + const next = [...rows] + next[index] = updated + return next +} + +function upsertRestServerInLive( + rows: ServerMetrics[] | undefined, + server: ServerResponse +): ServerMetrics[] | undefined { + if (!rows) { + return rows + } + const index = rows.findIndex((row) => row.id === server.id) + if (index < 0) { + return [...rows, projectRestServer(blankServerMetrics(server.id), server)] + } + const next = [...rows] + next[index] = projectRestServer(rows[index], server) + return next +} + +function upsertRestServerInList( + rows: ServerResponse[] | undefined, + server: ServerResponse +): ServerResponse[] | undefined { + if (!rows) { + return rows + } + const index = rows.findIndex((row) => row.id === server.id) + if (index < 0) { + return [...rows, server] + } + const next = [...rows] + next[index] = server + return next +} + +function updateExistingDetail( + queryClient: QueryClient, + serverId: string, + update: (server: ServerResponse) => ServerResponse +): void { + const key = serverDetailKey(serverId) + const current = queryClient.getQueryData(key) + if (current !== undefined) { + queryClient.setQueryData(key, update(current)) + } +} + +function removeMissingDetails(queryClient: QueryClient, serverIds: ReadonlySet): void { + const detailQueries = queryClient.getQueryCache().findAll({ queryKey: SERVER_DETAIL_PREFIX }) + for (const query of detailQueries) { + const serverId = query.queryKey[2] + if (typeof serverId === 'string' && !serverIds.has(serverId)) { + queryClient.removeQueries({ exact: true, queryKey: serverDetailKey(serverId) }) + } + } +} + +function setDockerAvailability(features: string[], available: boolean): string[] { + const hasDocker = features.includes('docker') + if (available === hasDocker) { + return features + } + return available ? [...features, 'docker'] : features.filter((feature) => feature !== 'docker') +} + +function assertNever(value: never): never { + throw new Error(`Unhandled server catalog event: ${JSON.stringify(value)}`) +} + +function ensureLiveProjectionLifetime(queryClient: QueryClient): void { + queryClient.setQueryDefaults(LIVE_SERVERS_KEY, { + gcTime: Number.POSITIVE_INFINITY, + staleTime: Number.POSITIVE_INFINITY + }) +} + +export function projectServerCatalog(queryClient: QueryClient, event: ServerCatalogEvent): void { + ensureLiveProjectionLifetime(queryClient) + switch (event.kind) { + case 'rest_snapshot': { + const serverIds = new Set(event.servers.map((server) => server.id)) + queryClient.setQueryData(LIVE_SERVERS_KEY, (current) => + projectRestSnapshot(current, event.servers) + ) + queryClient.setQueryData(SERVER_LIST_KEY, [...event.servers]) + removeMissingDetails(queryClient, serverIds) + for (const server of event.servers) { + updateExistingDetail(queryClient, server.id, () => server) + } + return + } + case 'server_saved': { + queryClient.setQueryData(LIVE_SERVERS_KEY, (current) => + upsertRestServerInLive(current, event.server) + ) + queryClient.setQueryData(SERVER_LIST_KEY, (current) => + upsertRestServerInList(current, event.server) + ) + queryClient.setQueryData(serverDetailKey(event.server.id), event.server) + return + } + case 'servers_removed': { + const removed = new Set(event.serverIds) + queryClient.setQueryData(LIVE_SERVERS_KEY, (current) => + current?.filter((server) => !removed.has(server.id)) + ) + queryClient.setQueryData(SERVER_LIST_KEY, (current) => + current?.filter((server) => !removed.has(server.id)) + ) + for (const serverId of event.serverIds) { + queryClient.removeQueries({ exact: true, queryKey: serverDetailKey(serverId) }) + } + return + } + case 'enrollment_changed': { + queryClient.setQueryData(LIVE_SERVERS_KEY, (current) => + updateById(current, event.serverId, (server) => ({ + ...server, + has_token: event.tokenRevoked ? false : server.has_token, + online: event.tokenRevoked ? false : server.online, + outstanding_enrollment: event.outstandingEnrollment + })) + ) + queryClient.setQueryData(SERVER_LIST_KEY, (current) => + updateById(current, event.serverId, (server) => ({ + ...server, + has_token: event.tokenRevoked ? false : server.has_token, + outstanding_enrollment: event.outstandingEnrollment + })) + ) + updateExistingDetail(queryClient, event.serverId, (server) => ({ + ...server, + has_token: event.tokenRevoked ? false : server.has_token, + outstanding_enrollment: event.outstandingEnrollment + })) + return + } + case 'tags_changed': { + queryClient.setQueryData(LIVE_SERVERS_KEY, (current) => + updateById(current, event.serverId, (server) => ({ ...server, tags: [...event.tags] })) + ) + return + } + case 'ws_full_sync': { + const serverIds = new Set(event.servers.map((server) => server.id)) + const fullSyncById = new Map(event.servers.map((server) => [server.id, server])) + const currentList = queryClient.getQueryData(SERVER_LIST_KEY) + const hasUnlistedServer = + currentList !== undefined && event.servers.some((server) => !currentList.some((row) => row.id === server.id)) + queryClient.setQueryData(LIVE_SERVERS_KEY, (current) => mergeWsFullSync(current, event.servers)) + queryClient.setQueryData(SERVER_LIST_KEY, (current) => + current + ?.filter((server) => serverIds.has(server.id)) + .map((server) => { + const fullSync = fullSyncById.get(server.id) + return fullSync ? projectFullSyncServerToRest(server, fullSync) : server + }) + ) + removeMissingDetails(queryClient, serverIds) + for (const server of event.servers) { + updateExistingDetail(queryClient, server.id, (current) => projectFullSyncServerToRest(current, server)) + } + if (hasUnlistedServer) { + queryClient.invalidateQueries({ exact: true, queryKey: SERVER_LIST_KEY }).catch(() => undefined) + } + return + } + case 'ws_update': { + queryClient.setQueryData(LIVE_SERVERS_KEY, (current) => + current ? mergeWsUpdate(current, event.servers) : [...event.servers] + ) + return + } + case 'online_changed': { + queryClient.setQueryData(LIVE_SERVERS_KEY, (current) => + updateById(current, event.serverId, (server) => ({ ...server, online: event.online })) + ) + return + } + case 'capabilities_changed': { + const agentLocalCapabilities = event.agentLocalCapabilities ?? null + const effectiveCapabilities = event.effectiveCapabilities ?? null + const temporary = event.temporary ? [...event.temporary] : [] + queryClient.setQueryData(LIVE_SERVERS_KEY, (current) => + updateById(current, event.serverId, (server) => ({ + ...server, + agent_local_capabilities: agentLocalCapabilities, + capabilities: event.capabilities, + effective_capabilities: effectiveCapabilities, + temporary + })) + ) + queryClient.setQueryData(SERVER_LIST_KEY, (current) => + updateById(current, event.serverId, (server) => ({ + ...server, + agent_local_capabilities: agentLocalCapabilities, + capabilities: event.capabilities, + effective_capabilities: effectiveCapabilities, + temporary + })) + ) + updateExistingDetail(queryClient, event.serverId, (server) => ({ + ...server, + agent_local_capabilities: agentLocalCapabilities, + capabilities: event.capabilities, + effective_capabilities: effectiveCapabilities, + temporary + })) + return + } + case 'agent_info_changed': { + const agentVersion = event.agentVersion ?? null + queryClient.setQueryData(LIVE_SERVERS_KEY, (current) => + updateById(current, event.serverId, (server) => ({ + ...server, + agent_version: agentVersion, + protocol_version: event.protocolVersion + })) + ) + queryClient.setQueryData(SERVER_LIST_KEY, (current) => + updateById(current, event.serverId, (server) => ({ + ...server, + agent_version: agentVersion, + protocol_version: event.protocolVersion + })) + ) + updateExistingDetail(queryClient, event.serverId, (server) => ({ + ...server, + agent_version: agentVersion, + protocol_version: event.protocolVersion + })) + return + } + case 'docker_availability_changed': { + queryClient.setQueryData(LIVE_SERVERS_KEY, (current) => + updateById(current, event.serverId, (server) => ({ + ...server, + features: setDockerAvailability(server.features ?? [], event.available) + })) + ) + queryClient.setQueryData(SERVER_LIST_KEY, (current) => + updateById(current, event.serverId, (server) => ({ + ...server, + features: setDockerAvailability(server.features, event.available) + })) + ) + updateExistingDetail(queryClient, event.serverId, (server) => ({ + ...server, + features: setDockerAvailability(server.features, event.available) + })) + return + } + default: + assertNever(event) + } +} + +export function useLiveServers(options: CatalogQueryOptions = {}) { + return useQuery({ + enabled: options.enabled ?? true, + gcTime: Number.POSITIVE_INFINITY, + queryFn: () => [], + queryKey: LIVE_SERVERS_KEY, + refetchOnMount: false, + refetchOnWindowFocus: false, + staleTime: Number.POSITIVE_INFINITY + }) +} + +export function useServerList(options: CatalogQueryOptions = {}) { + const queryClient = useQueryClient() + return useQuery({ + enabled: options.enabled ?? true, + queryFn: () => fetchAndProjectRestSnapshot(queryClient), + queryKey: SERVER_LIST_KEY + }) +} + +export function useServerDetail(serverId: string, options: CatalogQueryOptions = {}) { + const queryClient = useQueryClient() + return useQuery({ + enabled: (options.enabled ?? true) && serverId.length > 0, + queryFn: () => fetchAndProjectServerDetail(queryClient, serverId), + queryKey: serverDetailKey(serverId) + }) +} + +export async function refreshServerCatalog(queryClient: QueryClient): Promise { + const generation = nextListRequestGeneration(queryClient) + await queryClient.cancelQueries({ exact: true, queryKey: SERVER_LIST_KEY }) + if (!isCurrentListRequest(queryClient, generation)) { + return + } + try { + await queryClient.fetchQuery({ + queryFn: () => fetchAndProjectRestSnapshot(queryClient, generation), + queryKey: SERVER_LIST_KEY, + staleTime: 0 + }) + } catch (error) { + if (isCurrentListRequest(queryClient, generation)) { + throw error + } + } +} + +function nextListRequestGeneration(queryClient: QueryClient): number { + const generation = (listRequestGenerations.get(queryClient) ?? 0) + 1 + listRequestGenerations.set(queryClient, generation) + return generation +} + +function isCurrentListRequest(queryClient: QueryClient, generation: number): boolean { + return listRequestGenerations.get(queryClient) === generation +} + +async function fetchAndProjectRestSnapshot( + queryClient: QueryClient, + generation = nextListRequestGeneration(queryClient) +): Promise { + const servers = await api.get('/api/servers') + if (isCurrentListRequest(queryClient, generation)) { + projectServerCatalog(queryClient, { kind: 'rest_snapshot', servers }) + } + return servers +} + +async function fetchAndProjectServerDetail(queryClient: QueryClient, serverId: string): Promise { + const server = await api.get(`/api/servers/${serverId}`) + projectServerCatalog(queryClient, { kind: 'server_saved', server }) + return server +} + +export function invalidateServerDetail(queryClient: QueryClient, serverId: string): Promise { + return queryClient.invalidateQueries({ exact: true, queryKey: serverDetailKey(serverId) }) +} + +export function readLiveServers(queryClient: QueryClient): ServerMetrics[] | undefined { + return queryClient.getQueryData(LIVE_SERVERS_KEY) +} + +export function subscribeLiveServers( + queryClient: QueryClient, + listener: (servers: ServerMetrics[] | undefined) => void +): () => void { + return queryClient.getQueryCache().subscribe((event) => { + if (event.type !== 'added' && event.type !== 'removed' && event.type !== 'updated') { + return + } + const queryKey = event.query.queryKey + if ( + queryKey.length === LIVE_SERVERS_KEY.length && + queryKey[0] === LIVE_SERVERS_KEY[0] && + queryKey[1] === LIVE_SERVERS_KEY[1] + ) { + listener(readLiveServers(queryClient)) + } + }) +} diff --git a/apps/web/src/lib/widget-helpers.ts b/apps/web/src/lib/widget-helpers.ts index 2399936d8..a1e3190ca 100644 --- a/apps/web/src/lib/widget-helpers.ts +++ b/apps/web/src/lib/widget-helpers.ts @@ -1,6 +1,6 @@ import i18next from 'i18next' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import type { ServerMetricRecord, UptimeDailyEntry } from '@/lib/api-schema' +import type { ServerMetrics } from '@/lib/server-catalog' import { parseDiskIoJson } from './disk-io' import { activeLocale, formatDateTime } from './format' diff --git a/apps/web/src/routes/_authed/index.tsx b/apps/web/src/routes/_authed/index.tsx index 9e10b37dd..a6bee1775 100644 --- a/apps/web/src/routes/_authed/index.tsx +++ b/apps/web/src/routes/_authed/index.tsx @@ -1,11 +1,10 @@ -import { useQuery } from '@tanstack/react-query' import { createFileRoute } from '@tanstack/react-router' import { useMemo, useState } from 'react' import { DashboardEditorView } from '@/components/dashboard/dashboard-editor-view' import { useAuth } from '@/hooks/use-auth' import { useDashboard, useDashboards, useDefaultDashboard, useUpdateDashboard } from '@/hooks/use-dashboard' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { withMockServers } from '@/lib/dev-mock-servers' +import { useLiveServers } from '@/lib/server-catalog' export const Route = createFileRoute('/_authed/')({ component: DashboardPage @@ -15,13 +14,7 @@ export function DashboardPage() { const { user } = useAuth() const isAdmin = user?.role === 'admin' - const { data: rawServers = [] } = useQuery({ - queryKey: ['servers'], - queryFn: () => [], - staleTime: Number.POSITIVE_INFINITY, - refetchOnMount: false, - refetchOnWindowFocus: false - }) + const { data: rawServers = [] } = useLiveServers() const servers = useMemo(() => withMockServers(rawServers), [rawServers]) const { data: dashboards = [] } = useDashboards() diff --git a/apps/web/src/routes/_authed/ip-quality.tsx b/apps/web/src/routes/_authed/ip-quality.tsx index acc410ee8..8f1cc9b2d 100644 --- a/apps/web/src/routes/_authed/ip-quality.tsx +++ b/apps/web/src/routes/_authed/ip-quality.tsx @@ -1,26 +1,17 @@ -import { useQuery } from '@tanstack/react-query' import { createFileRoute } from '@tanstack/react-router' import { IpQualityContent } from '@/components/status/ip-quality-content' import { useIpQualityOverview, useIpQualityServices } from '@/hooks/use-ip-quality-api' -import { api } from '@/lib/api-client' +import { useServerList } from '@/lib/server-catalog' export const Route = createFileRoute('/_authed/ip-quality')({ component: IpQualityOverviewPage }) -interface ServerLite { - id: string - name: string -} - function IpQualityOverviewPage() { const { data: overview = [], isLoading: overviewLoading } = useIpQualityOverview() const { data: services = [], isLoading: servicesLoading } = useIpQualityServices() - const { data: servers = [], isLoading: serversLoading } = useQuery({ - queryKey: ['servers', 'lite'], - queryFn: () => api.get('/api/servers') - }) + const { data: servers = [], isLoading: serversLoading } = useServerList() const isLoading = overviewLoading || servicesLoading || serversLoading diff --git a/apps/web/src/routes/_authed/security/$serverId.tsx b/apps/web/src/routes/_authed/security/$serverId.tsx index 763dc66f9..48a2f22ec 100644 --- a/apps/web/src/routes/_authed/security/$serverId.tsx +++ b/apps/web/src/routes/_authed/security/$serverId.tsx @@ -1,4 +1,3 @@ -import { useQuery } from '@tanstack/react-query' import { createFileRoute, Link } from '@tanstack/react-router' import { useMemo, useState } from 'react' import { useTranslation } from 'react-i18next' @@ -8,8 +7,8 @@ import { SecurityKpiCards } from '@/components/security/kpi-cards' import { SecurityTimelineChart } from '@/components/security/timeline-chart' import { Button } from '@/components/ui/button' import { useSecurityEvents } from '@/hooks/use-security-events' -import { api } from '@/lib/api-client' import type { SecurityEventDto } from '@/lib/api-schema' +import { useServerDetail } from '@/lib/server-catalog' export const Route = createFileRoute('/_authed/security/$serverId')({ component: SecurityServerPage @@ -17,11 +16,6 @@ export const Route = createFileRoute('/_authed/security/$serverId')({ type RangeKey = '24h' | '7d' | '30d' -interface ServerLite { - id: string - name: string -} - const RANGE_HOURS: Record = { '24h': 24, '7d': 24 * 7, @@ -41,10 +35,7 @@ function SecurityServerPage() { const since = useMemo(() => computeSince(range), [range]) const eventsQuery = useSecurityEvents({ server_id: serverId, since, limit: 100 }) - const { data: server } = useQuery({ - queryKey: ['server', 'lite', serverId], - queryFn: () => api.get(`/api/servers/${serverId}`) - }) + const { data: server } = useServerDetail(serverId) const events = useMemo(() => { const out: SecurityEventDto[] = [] diff --git a/apps/web/src/routes/_authed/security/index.tsx b/apps/web/src/routes/_authed/security/index.tsx index ce1f539c2..a3bb8a66f 100644 --- a/apps/web/src/routes/_authed/security/index.tsx +++ b/apps/web/src/routes/_authed/security/index.tsx @@ -1,4 +1,3 @@ -import { useQuery } from '@tanstack/react-query' import { createFileRoute } from '@tanstack/react-router' import { useMemo, useReducer } from 'react' import { useTranslation } from 'react-i18next' @@ -12,8 +11,8 @@ import { Input } from '@/components/ui/input' import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from '@/components/ui/select' import { useAuth } from '@/hooks/use-auth' import { type SecurityEventFilters, useSecurityEvents } from '@/hooks/use-security-events' -import { api } from '@/lib/api-client' import type { SecurityEventDto } from '@/lib/api-schema' +import { useServerList } from '@/lib/server-catalog' export const Route = createFileRoute('/_authed/security/')({ component: SecurityIndexPage @@ -21,11 +20,6 @@ export const Route = createFileRoute('/_authed/security/')({ type RangeKey = '24h' | '7d' | '30d' -interface ServerSummary { - id: string - name: string -} - const RANGE_HOURS: Record = { '24h': 24, '7d': 24 * 7, @@ -119,10 +113,7 @@ function SecurityIndexPage() { const eventsQuery = useSecurityEvents(filters) - const { data: servers } = useQuery({ - queryKey: ['servers', 'lite'], - queryFn: () => api.get('/api/servers') - }) + const { data: servers } = useServerList() const allEvents = useMemo(() => { const list: SecurityEventDto[] = [] diff --git a/apps/web/src/routes/_authed/servers/$id-page.tsx b/apps/web/src/routes/_authed/servers/$id-page.tsx index 7bb26d5d6..00d26545d 100644 --- a/apps/web/src/routes/_authed/servers/$id-page.tsx +++ b/apps/web/src/routes/_authed/servers/$id-page.tsx @@ -16,10 +16,10 @@ import { Button } from '@/components/ui/button' import { Skeleton } from '@/components/ui/skeleton' import { useServer } from '@/hooks/use-api' import { useAuth } from '@/hooks/use-auth' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { api } from '@/lib/api-client' import type { ServerResponse } from '@/lib/api-schema' import { CAP_DOCKER, CAP_FILE, CAP_TERMINAL, getEffectiveCapabilityEnabled } from '@/lib/capabilities' +import { useLiveServers } from '@/lib/server-catalog' import { formatBytes } from '@/lib/utils' import { useUpgradeJobsStore } from '@/stores/upgrade-jobs-store' import type { ServerDetailTab } from './server-detail-search' @@ -163,13 +163,7 @@ export function ServerDetailPage() { const { data: server, isLoading: serverLoading } = useServer(id) - const { data: liveServers } = useQuery({ - queryKey: ['servers'], - queryFn: () => [], - staleTime: Number.POSITIVE_INFINITY, - refetchOnMount: false, - refetchOnWindowFocus: false - }) + const { data: liveServers } = useLiveServers() const liveHydrated = liveServers !== undefined const liveData = liveServers?.find((s) => s.id === id) const upgradeJob = useUpgradeJobsStore((state) => state.jobs.get(id)) diff --git a/apps/web/src/routes/_authed/servers/$serverId/docker/index.tsx b/apps/web/src/routes/_authed/servers/$serverId/docker/index.tsx index 1d700e128..f7d383280 100644 --- a/apps/web/src/routes/_authed/servers/$serverId/docker/index.tsx +++ b/apps/web/src/routes/_authed/servers/$serverId/docker/index.tsx @@ -4,9 +4,9 @@ import { ArrowLeft, Container, HardDrive, Network } from 'lucide-react' import { useState } from 'react' import { useTranslation } from 'react-i18next' import { Button } from '@/components/ui/button' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { api } from '@/lib/api-client' import { CAP_DOCKER, getEffectiveCapabilityEnabled } from '@/lib/capabilities' +import { useLiveServers, useServerDetail } from '@/lib/server-catalog' import { ContainerDetailDialog } from './components/container-detail-dialog' import { ContainerList } from './components/container-list' import { DockerEvents } from './components/docker-events' @@ -27,28 +27,12 @@ function DockerPage() { const [networksOpen, setNetworksOpen] = useState(false) const [volumesOpen, setVolumesOpen] = useState(false) - const { data: liveServers } = useQuery({ - queryKey: ['servers'], - queryFn: () => [], - staleTime: Number.POSITIVE_INFINITY, - refetchOnMount: false, - refetchOnWindowFocus: false - }) + const { data: liveServers } = useLiveServers() const liveServer = liveServers?.find((s) => s.id === serverId) const wsDockerAvailable = liveServer?.features?.includes('docker') ?? false // Server detail is the source of truth for capability gating and a fallback for features. - const { data: serverDetail } = useQuery({ - queryKey: ['servers', serverId], - queryFn: () => - api.get<{ - capabilities?: number - effective_capabilities?: number | null - features?: string[] - }>(`/api/servers/${serverId}`), - enabled: serverId.length > 0, - staleTime: 30_000 - }) + const { data: serverDetail } = useServerDetail(serverId) const dockerCapabilityEnabled = getEffectiveCapabilityEnabled( serverDetail?.effective_capabilities, diff --git a/apps/web/src/routes/_authed/servers/components/index-cells.tsx b/apps/web/src/routes/_authed/servers/components/index-cells.tsx index a576b9b8f..d3afc9976 100644 --- a/apps/web/src/routes/_authed/servers/components/index-cells.tsx +++ b/apps/web/src/routes/_authed/servers/components/index-cells.tsx @@ -7,8 +7,8 @@ import { MetricValue } from '@/components/server/metric-value' import { StatusDot } from '@/components/server/status-dot' import { deriveServerStatus } from '@/components/server/status-dot-utils' import { TagChipRow } from '@/components/server/tag-chip' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import type { TrafficOverviewItem } from '@/hooks/use-traffic-overview' +import type { ServerMetrics } from '@/lib/server-catalog' import { computeTrafficQuota } from '@/lib/traffic' import { cn, formatUptime } from '@/lib/utils' import { getBarColor, getBarTextColor } from './metric-bar-colors' diff --git a/apps/web/src/routes/_authed/servers/components/server-columns.tsx b/apps/web/src/routes/_authed/servers/components/server-columns.tsx index aea7c2631..e00ff0fdb 100644 --- a/apps/web/src/routes/_authed/servers/components/server-columns.tsx +++ b/apps/web/src/routes/_authed/servers/components/server-columns.tsx @@ -5,9 +5,9 @@ import { CostCell } from '@/components/server/cost-cell' import { StatusDot } from '@/components/server/status-dot' import { deriveServerStatus } from '@/components/server/status-dot-utils' import { Checkbox } from '@/components/ui/checkbox' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import type { TrafficOverviewItem } from '@/hooks/use-traffic-overview' import type { ServerCostOverview } from '@/lib/api-schema' +import type { ServerMetrics } from '@/lib/server-catalog' import { cn } from '@/lib/utils' import { CpuCell, DiskCell, MemoryCell, NameCell, NetworkCell, UptimeCell } from './index-cells' import { UpgradeBadgeCell } from './upgrade-badge-cell' diff --git a/apps/web/src/routes/_authed/servers/components/servers-page-toolbar.tsx b/apps/web/src/routes/_authed/servers/components/servers-page-toolbar.tsx index 485895b47..615e013cb 100644 --- a/apps/web/src/routes/_authed/servers/components/servers-page-toolbar.tsx +++ b/apps/web/src/routes/_authed/servers/components/servers-page-toolbar.tsx @@ -16,7 +16,7 @@ import { import { Button } from '@/components/ui/button' import { Input } from '@/components/ui/input' import { ToggleGroup, ToggleGroupItem } from '@/components/ui/toggle-group' -import type { ServerMetrics } from '@/hooks/use-servers-ws' +import type { ServerMetrics } from '@/lib/server-catalog' export function ServersPageToolbar({ batchDeletePending, diff --git a/apps/web/src/routes/_authed/servers/index.cells.test.tsx b/apps/web/src/routes/_authed/servers/index.cells.test.tsx index 55b04dd2a..2207a6034 100644 --- a/apps/web/src/routes/_authed/servers/index.cells.test.tsx +++ b/apps/web/src/routes/_authed/servers/index.cells.test.tsx @@ -1,7 +1,7 @@ import { render, screen } from '@testing-library/react' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import type { TrafficOverviewItem } from '@/hooks/use-traffic-overview' +import type { ServerMetrics } from '@/lib/server-catalog' import { CpuCell, DiskCell, diff --git a/apps/web/src/routes/_authed/servers/index.tsx b/apps/web/src/routes/_authed/servers/index.tsx index dc0e1553d..2ee0edfb4 100644 --- a/apps/web/src/routes/_authed/servers/index.tsx +++ b/apps/web/src/routes/_authed/servers/index.tsx @@ -15,12 +15,12 @@ import { useCostOverview } from '@/hooks/use-cost' import { useDataTable } from '@/hooks/use-data-table' import { useNetworkOverview, useNetworkSetting } from '@/hooks/use-network-api' import { useScrollViewportHeight } from '@/hooks/use-scroll-viewport-height' -import { reconcileServersFromRest, type ServerMetrics } from '@/hooks/use-servers-ws' import { useTrafficOverview } from '@/hooks/use-traffic-overview' import { api } from '@/lib/api-client' -import type { ServerGroup, ServerResponse } from '@/lib/api-schema' +import type { ServerGroup } from '@/lib/api-schema' import { withMockServers } from '@/lib/dev-mock-servers' import { countCleanupCandidates } from '@/lib/orphan-server-utils' +import { projectServerCatalog, refreshServerCatalog, type ServerMetrics, useLiveServers } from '@/lib/server-catalog' import { cn } from '@/lib/utils' import { getInitialServersView } from './components/mobile-view' import { buildServerColumns } from './components/server-columns' @@ -63,13 +63,7 @@ function ServersListPage() { navigate({ search: (prev) => ({ ...prev, view: value }) }) } - const { data: rawServers = [] } = useQuery({ - queryKey: ['servers'], - queryFn: () => [], - staleTime: Number.POSITIVE_INFINITY, - refetchOnMount: false, - refetchOnWindowFocus: false - }) + const { data: rawServers = [] } = useLiveServers() const servers = useMemo(() => withMockServers(rawServers), [rawServers]) const { data: groups } = useQuery({ @@ -168,13 +162,8 @@ function ServersListPage() { const cleanupMutation = useMutation({ mutationFn: () => api.delete<{ deleted_count: number }>('/api/servers/cleanup'), onSuccess: async (data) => { - // ['servers'] is a WS-fed cache (queryFn: () => []); refresh membership - // from REST instead of invalidating (which would wipe the visible list). try { - const fresh = await api.get('/api/servers') - queryClient.setQueryData(['servers'], (prev) => - reconcileServersFromRest(prev, fresh as unknown as Array & { id: string }>) - ) + await refreshServerCatalog(queryClient) } catch { // Best-effort: next WS full_sync will reconcile. } @@ -189,9 +178,7 @@ function ServersListPage() { mutationFn: (ids: string[]) => api.post<{ deleted: number }>('/api/servers/batch-delete', { ids }), onSuccess: (_data, ids) => { table.toggleAllRowsSelected(false) - const removed = new Set(ids) - // Same WS-cache caveat: filter the deleted ids out in place. - queryClient.setQueryData(['servers'], (prev) => prev?.filter((s) => !removed.has(s.id))) + projectServerCatalog(queryClient, { kind: 'servers_removed', serverIds: ids }) } }) diff --git a/apps/web/src/routes/_authed/settings/alerts.tsx b/apps/web/src/routes/_authed/settings/alerts.tsx index 742c3cc16..5965236f2 100644 --- a/apps/web/src/routes/_authed/settings/alerts.tsx +++ b/apps/web/src/routes/_authed/settings/alerts.tsx @@ -8,6 +8,7 @@ import { SecurityAlertPresets } from '@/components/security/alert-presets' import { Button } from '@/components/ui/button' import { api } from '@/lib/api-client' import type { AlertRule, AlertStateResponse, NotificationGroup } from '@/lib/api-schema' +import { useServerList } from '@/lib/server-catalog' import { AlertRuleForm, type CreateAlertRuleInput } from './alert-rule-form' import { AlertRulesList } from './alert-rules-list' @@ -15,11 +16,6 @@ export const Route = createFileRoute('/_authed/settings/alerts')({ component: AlertsPage }) -interface Server { - id: string - name: string -} - function AlertsPage() { const { t } = useTranslation(['settings', 'common']) const queryClient = useQueryClient() @@ -37,11 +33,7 @@ function AlertsPage() { queryFn: () => api.get('/api/notification-groups') }) - const { data: servers } = useQuery({ - queryKey: ['servers'], - queryFn: () => api.get('/api/servers'), - enabled: showForm - }) + const { data: servers } = useServerList({ enabled: showForm }) const { data: states } = useQuery({ queryKey: ['alert-rule-states', expandedRuleId], diff --git a/apps/web/src/routes/_authed/settings/capabilities.tsx b/apps/web/src/routes/_authed/settings/capabilities.tsx index 3e71a8cef..f1244126c 100644 --- a/apps/web/src/routes/_authed/settings/capabilities.tsx +++ b/apps/web/src/routes/_authed/settings/capabilities.tsx @@ -1,4 +1,3 @@ -import { useQuery } from '@tanstack/react-query' import { createFileRoute } from '@tanstack/react-router' import { type ColumnDef, getCoreRowModel, type Table, useReactTable } from '@tanstack/react-table' import { Check, Minus, Search, ShieldAlert } from 'lucide-react' @@ -8,7 +7,6 @@ import { TemporaryBadge } from '@/components/server/temporary-badge' import { DataTable } from '@/components/ui/data-table' import { Input } from '@/components/ui/input' import { Skeleton } from '@/components/ui/skeleton' -import type { ServerMetrics } from '@/hooks/use-servers-ws' import { CAPABILITIES, classifyCapability, @@ -16,6 +14,7 @@ import { RISK_TEXT_CLASS, temporaryGrantFor } from '@/lib/capabilities' +import { type ServerMetrics, useLiveServers } from '@/lib/server-catalog' export const Route = createFileRoute('/_authed/settings/capabilities')({ validateSearch: (search: Record) => ({ @@ -69,16 +68,8 @@ export function CapabilitiesPage() { const { q: search } = Route.useSearch() const navigate = Route.useNavigate() - // The ['servers'] cache is fed by the global WebSocket layer. Capabilities are - // agent-owned and reported by the agent over that channel, so this page is a - // read-only mirror of what each agent has enabled in its config file. - const { data: servers = [], isLoading } = useQuery({ - queryKey: ['servers'], - queryFn: () => [], - staleTime: Number.POSITIVE_INFINITY, - refetchOnMount: false, - refetchOnWindowFocus: false - }) + // Capabilities are agent-owned and reported over the global WebSocket. + const { data: servers = [], isLoading } = useLiveServers() const filtered = useMemo( () => servers.filter((s) => s.name.toLowerCase().includes(search.toLowerCase())), diff --git a/apps/web/src/routes/_authed/settings/ping-tasks.tsx b/apps/web/src/routes/_authed/settings/ping-tasks.tsx index 2f5e97f12..07ddbb6aa 100644 --- a/apps/web/src/routes/_authed/settings/ping-tasks.tsx +++ b/apps/web/src/routes/_authed/settings/ping-tasks.tsx @@ -19,9 +19,10 @@ import { Button } from '@/components/ui/button' import { Skeleton } from '@/components/ui/skeleton' import { api } from '@/lib/api-client' import type { PingTask } from '@/lib/api-schema' +import { useServerList } from '@/lib/server-catalog' import { PingResultsChart } from './ping-results-chart' import { PingTaskCreateDialog } from './ping-task-create-dialog' -import type { ProbeType, Server } from './ping-task-types' +import type { ProbeType } from './ping-task-types' export const Route = createFileRoute('/_authed/settings/ping-tasks')({ component: PingTasksPage @@ -44,10 +45,7 @@ function PingTasksPage() { queryFn: () => api.get('/api/ping-tasks') }) - const { data: servers } = useQuery({ - queryKey: ['servers-list'], - queryFn: () => api.get('/api/servers') - }) + const { data: servers } = useServerList() const deleteMutation = useMutation({ mutationFn: (id: string) => api.delete(`/api/ping-tasks/${id}`), diff --git a/apps/web/src/routes/_authed/settings/status-pages.tsx b/apps/web/src/routes/_authed/settings/status-pages.tsx index 4b56e00b9..51e55ee2d 100644 --- a/apps/web/src/routes/_authed/settings/status-pages.tsx +++ b/apps/web/src/routes/_authed/settings/status-pages.tsx @@ -1,9 +1,7 @@ -import { useQuery } from '@tanstack/react-query' import { createFileRoute } from '@tanstack/react-router' import { useTranslation } from 'react-i18next' import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs' -import { api } from '@/lib/api-client' -import type { ServerResponse } from '@/lib/api-schema' +import { useServerList } from '@/lib/server-catalog' import { StatusPageConfigForm } from './status-page-config-form' import { StatusPageIncidentsTab } from './status-page-incidents-tab' import { StatusPageMaintenanceTab } from './status-page-maintenance-tab' @@ -27,10 +25,7 @@ function StatusPagesManagement() { const { tab } = Route.useSearch() const navigate = Route.useNavigate() - const { data: servers } = useQuery({ - queryKey: ['servers-list'], - queryFn: () => api.get('/api/servers') - }) + const { data: servers } = useServerList() return (
diff --git a/apps/web/src/routes/_authed/settings/tasks.tsx b/apps/web/src/routes/_authed/settings/tasks.tsx index d18607c1a..20db99eec 100644 --- a/apps/web/src/routes/_authed/settings/tasks.tsx +++ b/apps/web/src/routes/_authed/settings/tasks.tsx @@ -14,18 +14,12 @@ import { api } from '@/lib/api-client' import type { TaskResponse, TaskResult } from '@/lib/api-schema' import { CAP_EXEC, getEffectiveCapabilityEnabled } from '@/lib/capabilities' import { formatDateTime } from '@/lib/format' +import { useServerList } from '@/lib/server-catalog' export const Route = createFileRoute('/_authed/settings/tasks')({ component: TasksPage }) -interface ServerInfo { - capabilities?: number - effective_capabilities?: number | null - id: string - name: string -} - function TasksPage() { const { t } = useTranslation(['settings', 'common']) @@ -55,10 +49,7 @@ function OneshotTaskPanel() { const [timeout, setTimeout] = useState(30) const [expandedTask, setExpandedTask] = useState(null) - const { data: servers } = useQuery({ - queryKey: ['servers-list'], - queryFn: () => api.get('/api/servers') - }) + const { data: servers } = useServerList() const createMutation = useMutation({ mutationFn: (input: { command: string; server_ids: string[]; timeout: number }) => diff --git a/apps/web/src/widgets-runtime/runtime-bridge.test.ts b/apps/web/src/widgets-runtime/runtime-bridge.test.ts index 8d49cc865..90ab1e56f 100644 --- a/apps/web/src/widgets-runtime/runtime-bridge.test.ts +++ b/apps/web/src/widgets-runtime/runtime-bridge.test.ts @@ -2,7 +2,7 @@ import * as Sdk from '@serverbee/widget-sdk' import { QueryClient } from '@tanstack/react-query' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' -import type { ServerMetrics } from '@/hooks/use-servers-ws' +import { projectServerCatalog, type ServerMetrics } from '@/lib/server-catalog' import { mountRuntimeBridge } from './runtime-bridge' function makeServer(id: string): ServerMetrics { @@ -41,6 +41,10 @@ function makeServer(id: string): ServerMetrics { } } +function setLiveServers(queryClient: QueryClient, servers: ServerMetrics[]): void { + projectServerCatalog(queryClient, { kind: 'ws_full_sync', servers }) +} + describe('runtime-bridge: React Query store wiring', () => { let qc: QueryClient @@ -59,22 +63,22 @@ describe('runtime-bridge: React Query store wiring', () => { }) it('serversStore reflects current cache after setQueryData', () => { - qc.setQueryData(['servers'], [makeServer('s1'), makeServer('s2')]) + setLiveServers(qc, [makeServer('s1'), makeServer('s2')]) const list = Sdk.getRuntime().serversStore() expect(list.map((s) => s.id)).toEqual(['s1', 's2']) }) it('serversStore returns the same array reference when cache is unchanged', () => { - qc.setQueryData(['servers'], [makeServer('s1')]) + setLiveServers(qc, [makeServer('s1')]) const a = Sdk.getRuntime().serversStore() const b = Sdk.getRuntime().serversStore() expect(a).toBe(b) }) - it('subscribeServers fires on ["servers"] cache updates', () => { + it('subscribeServers fires on live catalog updates', () => { const cb = vi.fn() const unsub = Sdk.getRuntime().subscribeServers(cb) - qc.setQueryData(['servers'], [makeServer('s1')]) + setLiveServers(qc, [makeServer('s1')]) expect(cb).toHaveBeenCalled() unsub() }) @@ -88,7 +92,7 @@ describe('runtime-bridge: React Query store wiring', () => { }) it('serverByIdStore looks up by id from cache', () => { - qc.setQueryData(['servers'], [makeServer('s1'), makeServer('s2')]) + setLiveServers(qc, [makeServer('s1'), makeServer('s2')]) const lookup = Sdk.getRuntime().serverByIdStore('s2') expect((lookup as ServerMetrics).id).toBe('s2') expect(Sdk.getRuntime().serverByIdStore('missing')).toBeUndefined() diff --git a/apps/web/src/widgets-runtime/runtime-bridge.ts b/apps/web/src/widgets-runtime/runtime-bridge.ts index e4d420c0a..63d1c62cb 100644 --- a/apps/web/src/widgets-runtime/runtime-bridge.ts +++ b/apps/web/src/widgets-runtime/runtime-bridge.ts @@ -12,9 +12,7 @@ import * as JsxRuntime from 'react/jsx-runtime' // biome-ignore lint/performance/noNamespaceImport: shim requires full namespace import * as ReactDOM from 'react-dom' import { toast } from 'sonner' -import type { ServerMetrics } from '@/hooks/use-servers-ws' - -const SERVERS_QUERY_KEY = ['servers'] as const +import { readLiveServers, type ServerMetrics, subscribeLiveServers } from '@/lib/server-catalog' export interface BridgeInputs { onConfigUpdate?: (instanceId: string, patch: Record) => void @@ -47,7 +45,7 @@ interface ServerSummariesCache { function makeServerSummariesGetter(queryClient: QueryClient): () => ServerSummary[] { let cache: ServerSummariesCache = { raw: undefined, summaries: [] } return () => { - const raw = queryClient.getQueryData(SERVERS_QUERY_KEY) + const raw = readLiveServers(queryClient) if (raw === cache.raw) { return cache.summaries } @@ -59,24 +57,13 @@ function makeServerSummariesGetter(queryClient: QueryClient): () => ServerSummar function makeServerByIdGetter(queryClient: QueryClient): (id: string) => unknown { return (id: string) => { - const raw = queryClient.getQueryData(SERVERS_QUERY_KEY) + const raw = readLiveServers(queryClient) return raw?.find((s) => s.id === id) } } function makeSubscribeServers(queryClient: QueryClient): (cb: () => void) => () => void { - return (cb: () => void) => { - const unsub = queryClient.getQueryCache().subscribe((event) => { - // We only care about updates to the ['servers'] cache. The cache emits - // a flurry of events (added/removed/updated/observers); filter to those - // that mutate the data for our queryKey. - const key = event.query.queryKey - if (Array.isArray(key) && key.length >= 1 && key[0] === SERVERS_QUERY_KEY[0]) { - cb() - } - }) - return unsub - } + return (cb: () => void) => subscribeLiveServers(queryClient, cb) } interface ThemeWatcher { From b9f611ab3437fd2041199c3640be338745ea8645 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 13:44:30 +0800 Subject: [PATCH 07/44] refactor(server): own agent desired-state sync in one reconciler Connect and mutation paths each carried their own fetch/map/send logic for ping tasks, network probes, IP-quality services, and the firewall blocklist, so capability filtering and protocol gating were duplicated and could drift. AgentDesiredStateReconciler now owns the full fetch -> map -> send lifecycle per domain, with per-provider locks so a newer snapshot can never be overwritten by an older concurrent one. - SystemInfo handshake reconciles every domain from live capabilities instead of hand-patching only the firewall - capabilities_changed drives a targeted per-domain reconcile; revoking ping/firewall caps now clears the corresponding agent state - ping/network/IP-quality/firewall CRUD routes signal the reconciler instead of hand-rolling their own sync sends - firewall reset ack only clears target state on ok=true, keeping the desired target when the agent reports a failed reset - PingService keeps DB/CRUD semantics only; AgentManager sync moves out --- crates/server/src/router/api/ip_quality.rs | 75 +- crates/server/src/router/api/network_probe.rs | 75 +- crates/server/src/router/api/ping.rs | 20 +- crates/server/src/router/api/server.rs | 28 +- crates/server/src/router/ws/agent/mod.rs | 61 -- crates/server/src/router/ws/agent/security.rs | 21 + .../server/src/router/ws/agent/system_info.rs | 35 +- crates/server/src/service/agent_reconcile.rs | 687 ++++++++++++++++++ crates/server/src/service/firewall.rs | 83 ++- crates/server/src/service/mod.rs | 1 + crates/server/src/service/ping.rs | 523 +------------ crates/server/src/state.rs | 8 +- crates/server/tests/agent_messages.rs | 106 ++- crates/server/tests/integration.rs | 12 +- crates/server/tests/router_probes_extra.rs | 67 +- 15 files changed, 1054 insertions(+), 748 deletions(-) create mode 100644 crates/server/src/service/agent_reconcile.rs diff --git a/crates/server/src/router/api/ip_quality.rs b/crates/server/src/router/api/ip_quality.rs index 63a831d37..ef10a4ab6 100644 --- a/crates/server/src/router/api/ip_quality.rs +++ b/crates/server/src/router/api/ip_quality.rs @@ -6,6 +6,7 @@ use axum::{Json, Router}; use serde::Deserialize; use crate::error::{ApiResponse, AppError, ok}; +use crate::service::agent_reconcile::AgentDesiredStateDomain; use crate::service::ip_quality::{ CreateCustomServiceInput, IpQualityService, IpQualitySettingDto, ServerIpQualityData, UnlockEventDto, UpdateServiceInput, @@ -41,35 +42,6 @@ pub fn write_router() -> Router> { .route("/ip-quality/servers/{id}/check", post(check_server)) } -// --------------------------------------------------------------------------- -// IpQualitySync re-broadcast helper -// --------------------------------------------------------------------------- - -/// Re-send `IpQualitySync` to every currently-online agent. -/// -/// Spec §4 requires `IpQualitySync` to be pushed on connect, on catalog -/// change, and on settings change. The WS handler covers the connect case; -/// this helper covers catalog/settings mutations so a change reaches already -/// connected agents without waiting for them to reconnect. Mirrors how -/// `network_probe.rs` re-broadcasts `NetworkProbeSync` after a mutation. -async fn broadcast_ip_quality_sync(state: &Arc) -> Result<(), AppError> { - let services = IpQualityService::enabled_service_defs(&state.db).await?; - let setting = IpQualityService::get_setting(&state.db).await?; - - for server_id in state.agent_manager.connected_server_ids() { - if let Some(tx) = state.agent_manager.get_sender(&server_id) { - let _ = tx - .send(ServerMessage::IpQualitySync { - services: services.clone(), - interval_hours: setting.check_interval_hours as u32, - }) - .await; - } - } - - Ok(()) -} - // --------------------------------------------------------------------------- // Query params // --------------------------------------------------------------------------- @@ -196,9 +168,7 @@ async fn create_service( Json(input): Json, ) -> Result>, AppError> { let service = IpQualityService::create_custom_service(&state.db, input).await?; - if let Err(e) = broadcast_ip_quality_sync(&state).await { - tracing::warn!("IpQualitySync broadcast failed: {e}"); - } + reconcile_ip_quality(&state).await; ok(service) } @@ -220,9 +190,7 @@ async fn update_service( Json(input): Json, ) -> Result>, AppError> { let service = IpQualityService::update_service(&state.db, &id, input).await?; - if let Err(e) = broadcast_ip_quality_sync(&state).await { - tracing::warn!("IpQualitySync broadcast failed: {e}"); - } + reconcile_ip_quality(&state).await; ok(service) } @@ -243,9 +211,7 @@ async fn delete_service( Path(id): Path, ) -> Result>, AppError> { IpQualityService::delete_service(&state.db, &id).await?; - if let Err(e) = broadcast_ip_quality_sync(&state).await { - tracing::warn!("IpQualitySync broadcast failed: {e}"); - } + reconcile_ip_quality(&state).await; ok("ok") } @@ -265,14 +231,21 @@ async fn update_settings( State(state): State>, Json(input): Json, ) -> Result>, AppError> { - let setting = - IpQualityService::update_setting(&state.db, input.check_interval_hours).await?; - if let Err(e) = broadcast_ip_quality_sync(&state).await { - tracing::warn!("IpQualitySync broadcast failed: {e}"); - } + let setting = IpQualityService::update_setting(&state.db, input.check_interval_hours).await?; + reconcile_ip_quality(&state).await; ok(setting) } +async fn reconcile_ip_quality(state: &AppState) { + if let Err(error) = state + .agent_desired_state + .reconcile_connected(AgentDesiredStateDomain::IpQuality) + .await + { + tracing::warn!(%error, "failed to reconcile IP quality desired state"); + } +} + #[utoipa::path( post, path = "/api/ip-quality/servers/{id}/check", @@ -337,7 +310,7 @@ mod tests { } #[tokio::test] - async fn broadcast_ip_quality_sync_reaches_online_agents() { + async fn reconcile_ip_quality_reaches_online_agents() { let (db, _tmp) = setup_test_db().await; let state = AppState::new(db, AppConfig::default()).await.unwrap(); @@ -347,7 +320,11 @@ mod tests { .agent_manager .add_connection("srv-online".into(), "Online".into(), tx, test_addr()); - broadcast_ip_quality_sync(&state).await.unwrap(); + state + .agent_desired_state + .reconcile_connected(AgentDesiredStateDomain::IpQuality) + .await + .unwrap(); // The online agent should receive an IpQualitySync with the 9 seeded // built-in services and the default 12h interval. @@ -369,12 +346,16 @@ mod tests { } #[tokio::test] - async fn broadcast_ip_quality_sync_with_no_agents_is_noop() { + async fn reconcile_ip_quality_with_no_agents_is_noop() { let (db, _tmp) = setup_test_db().await; let state = AppState::new(db, AppConfig::default()).await.unwrap(); // No connected agents — must succeed without error. - broadcast_ip_quality_sync(&state).await.unwrap(); + state + .agent_desired_state + .reconcile_connected(AgentDesiredStateDomain::IpQuality) + .await + .unwrap(); } // ----------------------------------------------------------------------- diff --git a/crates/server/src/router/api/network_probe.rs b/crates/server/src/router/api/network_probe.rs index 03bf760f5..79f1521f2 100644 --- a/crates/server/src/router/api/network_probe.rs +++ b/crates/server/src/router/api/network_probe.rs @@ -6,13 +6,12 @@ use axum::{Json, Router}; use serde::Deserialize; use crate::error::{ApiResponse, AppError, ok}; +use crate::service::agent_reconcile::AgentDesiredStateDomain; use crate::service::network_probe::{ CreateNetworkProbeTarget, NetworkProbeAnomaly, NetworkProbeService, NetworkProbeSetting, ServerOverview, TargetDto, UpdateNetworkProbeTarget, }; use crate::state::AppState; -use serverbee_common::protocol::ServerMessage; -use serverbee_common::types::NetworkProbeTarget; /// GET endpoints accessible to all authenticated users (admin + member). pub fn read_router() -> Router> { @@ -129,6 +128,7 @@ async fn update_target( Json(input): Json, ) -> Result>, AppError> { let target = NetworkProbeService::update_target(&state.db, &id, input).await?; + reconcile_network_probes(&state).await; ok(target) } @@ -147,45 +147,9 @@ async fn delete_target( State(state): State>, Path(id): Path, ) -> Result>, AppError> { - // Find which servers have this target configured, before deletion - use crate::entity::network_probe_config; - use sea_orm::ColumnTrait; - use sea_orm::EntityTrait; - use sea_orm::QueryFilter; - - let affected_configs = network_probe_config::Entity::find() - .filter(network_probe_config::Column::TargetId.eq(id.as_str())) - .all(&state.db) - .await?; - let affected_server_ids: Vec = - affected_configs.into_iter().map(|c| c.server_id).collect(); - // Delete the target (cascades records + configs + setting cleanup) NetworkProbeService::delete_target(&state.db, &id).await?; - - // Notify affected agents - let setting = NetworkProbeService::get_setting(&state.db).await?; - for server_id in &affected_server_ids { - let targets = NetworkProbeService::get_server_targets(&state.db, server_id).await?; - let probe_targets: Vec = targets - .into_iter() - .map(|t| NetworkProbeTarget { - target_id: t.id, - name: t.name, - target: t.target, - probe_type: t.probe_type, - }) - .collect(); - - if let Some(tx) = state.agent_manager.get_sender(server_id) { - let msg = ServerMessage::NetworkProbeSync { - targets: probe_targets, - interval: setting.interval, - packet_count: setting.packet_count, - }; - let _ = tx.send(msg).await; - } - } + reconcile_network_probes(&state).await; ok("ok") } @@ -205,32 +169,19 @@ async fn update_setting( Json(input): Json, ) -> Result>, AppError> { NetworkProbeService::update_setting(&state.db, &input).await?; + reconcile_network_probes(&state).await; - // Push updated config to all currently-online agents - let online_ids = state.agent_manager.connected_server_ids(); - for server_id in &online_ids { - let targets = NetworkProbeService::get_server_targets(&state.db, server_id).await?; - let probe_targets: Vec = targets - .into_iter() - .map(|t| NetworkProbeTarget { - target_id: t.id, - name: t.name, - target: t.target, - probe_type: t.probe_type, - }) - .collect(); + ok(input) +} - if let Some(tx) = state.agent_manager.get_sender(server_id) { - let msg = ServerMessage::NetworkProbeSync { - targets: probe_targets, - interval: input.interval, - packet_count: input.packet_count, - }; - let _ = tx.send(msg).await; - } +async fn reconcile_network_probes(state: &AppState) { + if let Err(error) = state + .agent_desired_state + .reconcile_connected(AgentDesiredStateDomain::NetworkProbes) + .await + { + tracing::warn!(%error, "failed to reconcile network probe desired state"); } - - ok(input) } // --------------------------------------------------------------------------- diff --git a/crates/server/src/router/api/ping.rs b/crates/server/src/router/api/ping.rs index 7e136b269..b4a76bd7a 100644 --- a/crates/server/src/router/api/ping.rs +++ b/crates/server/src/router/api/ping.rs @@ -8,6 +8,7 @@ use serde::Deserialize; use crate::entity::{ping_record, ping_task}; use crate::error::{ApiResponse, AppError, ok}; +use crate::service::agent_reconcile::AgentDesiredStateDomain; use crate::service::ping::{CreatePingTask, PingService, UpdatePingTask}; use crate::state::AppState; @@ -80,7 +81,8 @@ async fn create_task( State(state): State>, Json(input): Json, ) -> Result>, AppError> { - let task = PingService::create(&state.db, &state.agent_manager, input).await?; + let task = PingService::create(&state.db, input).await?; + reconcile_ping_tasks(&state).await; ok(task) } @@ -102,7 +104,8 @@ async fn update_task( Path(id): Path, Json(input): Json, ) -> Result>, AppError> { - let task = PingService::update(&state.db, &state.agent_manager, &id, input).await?; + let task = PingService::update(&state.db, &id, input).await?; + reconcile_ping_tasks(&state).await; ok(task) } @@ -122,10 +125,21 @@ async fn delete_task( State(state): State>, Path(id): Path, ) -> Result>, AppError> { - PingService::delete(&state.db, &state.agent_manager, &id).await?; + PingService::delete(&state.db, &id).await?; + reconcile_ping_tasks(&state).await; ok("ok") } +async fn reconcile_ping_tasks(state: &AppState) { + if let Err(error) = state + .agent_desired_state + .reconcile_connected(AgentDesiredStateDomain::PingTasks) + .await + { + tracing::warn!(%error, "failed to reconcile ping task desired state"); + } +} + #[derive(Deserialize, utoipa::IntoParams)] pub struct RecordsQuery { from: DateTime, diff --git a/crates/server/src/router/api/server.rs b/crates/server/src/router/api/server.rs index 159817e46..b0837d21e 100644 --- a/crates/server/src/router/api/server.rs +++ b/crates/server/src/router/api/server.rs @@ -23,6 +23,7 @@ use crate::router::api::network_probe::{ get_server_network_targets, }; use crate::service::agent_manager::AgentManager; +use crate::service::agent_reconcile::AgentDesiredStateDomain; use crate::service::audit::AuditService; use crate::service::enrollment::{DEFAULT_TTL_SECS, EnrollmentService}; use crate::service::network_probe::NetworkProbeService; @@ -33,7 +34,7 @@ use crate::service::task_scheduler; use crate::service::upgrade_tracker::{StartUpgradeJobError, UpgradeLookup}; use crate::state::AppState; use serverbee_common::protocol::ServerMessage; -use serverbee_common::types::{NetworkProbeTarget, OutstandingEnrollmentSummary}; +use serverbee_common::types::OutstandingEnrollmentSummary; const DEFAULT_SERVER_NAME: &str = "New Server"; @@ -1198,25 +1199,12 @@ async fn set_server_network_targets( ) -> Result>, AppError> { NetworkProbeService::set_server_targets(&state.db, &id, body.target_ids).await?; - // Push updated NetworkProbeSync to agent if online - if let Some(tx) = state.agent_manager.get_sender(&id) { - let targets = NetworkProbeService::get_server_targets(&state.db, &id).await?; - let setting = NetworkProbeService::get_setting(&state.db).await?; - let probe_targets: Vec = targets - .into_iter() - .map(|t| NetworkProbeTarget { - target_id: t.id, - name: t.name, - target: t.target, - probe_type: t.probe_type, - }) - .collect(); - let msg = ServerMessage::NetworkProbeSync { - targets: probe_targets, - interval: setting.interval, - packet_count: setting.packet_count, - }; - let _ = tx.send(msg).await; + if let Err(error) = state + .agent_desired_state + .reconcile_agent(&id, AgentDesiredStateDomain::NetworkProbes) + .await + { + tracing::warn!(server_id = id, %error, "failed to reconcile network probe desired state"); } ok("ok") diff --git a/crates/server/src/router/ws/agent/mod.rs b/crates/server/src/router/ws/agent/mod.rs index 4d314f2a1..b4f1e2ec2 100644 --- a/crates/server/src/router/ws/agent/mod.rs +++ b/crates/server/src/router/ws/agent/mod.rs @@ -14,15 +14,11 @@ use tokio::sync::mpsc; use crate::router::utils::extract_client_ip; use crate::service::auth::AuthService; -use crate::service::ip_quality::IpQualityService; -use crate::service::network_probe::NetworkProbeService; -use crate::service::ping::PingService; use crate::service::record::RecordService; use crate::service::upgrade_tracker::UpgradeLookup; use crate::state::AppState; use serverbee_common::constants::MAX_WS_MESSAGE_SIZE; use serverbee_common::protocol::{AgentMessage, ServerMessage}; -use serverbee_common::types::NetworkProbeTarget as NetworkProbeTargetDto; mod docker; mod file_transfer; @@ -182,63 +178,6 @@ async fn handle_agent_ws( connection_id }; - // Send current ping tasks to the newly connected agent - PingService::sync_tasks_to_agent(&state.db, &state.agent_manager, &server_id).await; - - // Send network probe sync to the newly connected agent - match NetworkProbeService::get_server_targets(&state.db, &server_id).await { - Ok(targets) => match NetworkProbeService::get_setting(&state.db).await { - Ok(setting) => { - let target_dtos: Vec = targets - .into_iter() - .map(|t| NetworkProbeTargetDto { - target_id: t.id, - name: t.name, - target: t.target, - probe_type: t.probe_type, - }) - .collect(); - if let Some(tx) = state.agent_manager.get_sender(&server_id) { - let _ = tx - .send(ServerMessage::NetworkProbeSync { - targets: target_dtos, - interval: setting.interval, - packet_count: setting.packet_count, - }) - .await; - } - } - Err(e) => { - tracing::error!("Failed to get network probe setting for {server_id}: {e}"); - } - }, - Err(e) => { - tracing::error!("Failed to get network probe targets for {server_id}: {e}"); - } - } - - // Send IP quality sync to the newly connected agent (mirrors NetworkProbeSync) - match IpQualityService::enabled_service_defs(&state.db).await { - Ok(services) => match IpQualityService::get_setting(&state.db).await { - Ok(setting) => { - if let Some(tx) = state.agent_manager.get_sender(&server_id) { - let _ = tx - .send(ServerMessage::IpQualitySync { - services, - interval_hours: setting.check_interval_hours as u32, - }) - .await; - } - } - Err(e) => { - tracing::error!("Failed to get IP quality setting for {server_id}: {e}"); - } - }, - Err(e) => { - tracing::error!("Failed to get IP quality service defs for {server_id}: {e}"); - } - } - tracing::info!("Agent {server_id} connected from {remote_addr}"); // Spawn a task to forward mpsc messages to WebSocket + send periodic Pings diff --git a/crates/server/src/router/ws/agent/security.rs b/crates/server/src/router/ws/agent/security.rs index fe2eede63..dcedc69cd 100644 --- a/crates/server/src/router/ws/agent/security.rs +++ b/crates/server/src/router/ws/agent/security.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use std::time::Duration; +use crate::service::agent_reconcile::AgentDesiredStateDomain; use crate::service::alert::AlertService; use crate::service::audit::AuditService; use crate::service::ip_quality::IpQualityService; @@ -186,6 +187,26 @@ pub(super) async fn on_capabilities_changed( tracing::error!("Failed to mirror capabilities for {server_id}: {e}"); } + for domain in [ + AgentDesiredStateDomain::PingTasks, + AgentDesiredStateDomain::NetworkProbes, + AgentDesiredStateDomain::IpQuality, + AgentDesiredStateDomain::Firewall, + ] { + if let Err(error) = state + .agent_desired_state + .reconcile_agent(server_id, domain) + .await + { + tracing::warn!( + server_id, + ?domain, + error = %error, + "capability-change desired-state reconcile failed" + ); + } + } + // Resolve display name + originating IP for the audit trail. Neither // `server_name` nor `remote_addr` is in scope here, so we look them // up from the DB / connection registry (mirroring the SystemInfo arm). diff --git a/crates/server/src/router/ws/agent/system_info.rs b/crates/server/src/router/ws/agent/system_info.rs index 17a19e626..299ac9132 100644 --- a/crates/server/src/router/ws/agent/system_info.rs +++ b/crates/server/src/router/ws/agent/system_info.rs @@ -1,5 +1,6 @@ //! `SystemInfo` / `IpChanged` handling: agent identity, GeoIP resolution, -//! capability mirroring, IP-change detection, and reconnect-time firewall sync. +//! capability mirroring, IP-change detection, and connection-time desired-state +//! reconciliation. use std::sync::Arc; @@ -233,30 +234,16 @@ pub(super) async fn on_system_info( } } - // Firewall blocklist: on every fresh SystemInfo, drop whatever the - // agent may have leftover from a previous boot, then resend the - // authoritative set. Gated on capability + protocol version. + if let Err(error) = state + .agent_desired_state + .reconcile_connection(server_id) + .await { - use serverbee_common::constants::{CAP_FIREWALL_BLOCK, has_capability}; - use serverbee_common::firewall::FIREWALL_MIN_PROTOCOL; - - let caps = state - .agent_manager - .get_effective_capabilities(server_id) - .unwrap_or(0); - if has_capability(caps, CAP_FIREWALL_BLOCK) && agent_pv >= FIREWALL_MIN_PROTOCOL { - state - .firewall - .push_reset_to(server_id, &state.agent_manager) - .await; - if let Err(e) = state - .firewall - .push_sync_to(server_id, &state.agent_manager) - .await - { - tracing::warn!(server_id, error = %e, "firewall sync push failed"); - } - } + tracing::warn!( + server_id, + error = %error, + "connection desired-state reconcile was incomplete" + ); } } diff --git a/crates/server/src/service/agent_reconcile.rs b/crates/server/src/service/agent_reconcile.rs new file mode 100644 index 000000000..1798bc8be --- /dev/null +++ b/crates/server/src/service/agent_reconcile.rs @@ -0,0 +1,687 @@ +//! Reconciles the server-owned desired state that agents execute. +//! +//! Routers and WebSocket handlers own transport concerns. This module owns the +//! fetch -> map -> send lifecycle for every full-state agent configuration, +//! including provider-specific capability and protocol policy. Firewall +//! mutations deliberately remain incremental; +//! the firewall provider here is the authoritative reset + sync path used on +//! connection and for explicit full repairs. + +use std::collections::HashMap; +use std::sync::Arc; + +use sea_orm::{ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter}; +use serverbee_common::constants::{ + CAP_DEFAULT, CAP_FIREWALL_BLOCK, has_capability, probe_type_to_cap, +}; +use serverbee_common::firewall::FIREWALL_MIN_PROTOCOL; +use serverbee_common::protocol::ServerMessage; +use serverbee_common::types::{NetworkProbeTarget, PingTaskConfig}; +use tokio::sync::Mutex; + +use crate::entity::{ping_task, server}; +use crate::error::AppError; +use crate::service::agent_manager::AgentManager; +use crate::service::firewall::FirewallService; +use crate::service::ip_quality::IpQualityService; +use crate::service::network_probe::{NetworkProbeService, NetworkProbeSetting, TargetDto}; + +/// A full-state configuration owned by the server and executed by an agent. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum AgentDesiredStateDomain { + PingTasks, + NetworkProbes, + IpQuality, + Firewall, +} + +/// Serializes each provider's read-and-send sequence. +/// +/// A mutation commits before it requests reconciliation. Holding the provider +/// lock across both the database read and the sends means two concurrent +/// mutations can send the same newest snapshot, but can never send an older +/// snapshot after a newer one. +#[derive(Default)] +struct ProviderLocks { + ping_tasks: Mutex<()>, + network_probes: Mutex<()>, + ip_quality: Mutex<()>, + firewall: Mutex<()>, +} + +/// Owns projection of persisted desired state into agent protocol messages. +#[derive(Clone)] +pub struct AgentDesiredStateReconciler { + db: DatabaseConnection, + agent_manager: Arc, + firewall: Arc, + locks: Arc, +} + +impl AgentDesiredStateReconciler { + pub fn new( + db: DatabaseConnection, + agent_manager: Arc, + firewall: Arc, + ) -> Self { + Self { + db, + agent_manager, + firewall, + locks: Arc::new(ProviderLocks::default()), + } + } + + /// Reconcile every desired-state domain after a new connection reports its + /// live capabilities and protocol version. All domains are attempted even + /// if one provider fails. + pub async fn reconcile_connection(&self, server_id: &str) -> Result<(), AppError> { + let mut first_error = None; + for domain in [ + AgentDesiredStateDomain::PingTasks, + AgentDesiredStateDomain::NetworkProbes, + AgentDesiredStateDomain::IpQuality, + AgentDesiredStateDomain::Firewall, + ] { + if let Err(error) = self.reconcile_agent(server_id, domain).await { + tracing::warn!( + server_id, + ?domain, + error = %error, + "failed to reconcile agent desired state" + ); + if first_error.is_none() { + first_error = Some(error); + } + } + } + + match first_error { + Some(error) => Err(error), + None => Ok(()), + } + } + + /// Reconcile one provider for one agent. + pub async fn reconcile_agent( + &self, + server_id: &str, + domain: AgentDesiredStateDomain, + ) -> Result<(), AppError> { + match domain { + AgentDesiredStateDomain::PingTasks => { + let _guard = self.locks.ping_tasks.lock().await; + let tasks = self.load_ping_tasks().await?; + let capabilities = self.capabilities_for_agent(server_id).await?; + self.send_ping_tasks(server_id, &tasks, capabilities).await; + Ok(()) + } + AgentDesiredStateDomain::NetworkProbes => { + let _guard = self.locks.network_probes.lock().await; + let setting = NetworkProbeService::get_setting(&self.db).await?; + self.send_network_probes(server_id, &setting).await + } + AgentDesiredStateDomain::IpQuality => { + let _guard = self.locks.ip_quality.lock().await; + self.send_ip_quality(server_id).await + } + AgentDesiredStateDomain::Firewall => { + let _guard = self.locks.firewall.lock().await; + self.send_firewall(server_id).await + } + } + } + + /// Reconcile one provider for every agent that is online at the start of + /// this call. Provider-global state is fetched once and reused for the + /// complete fan-out. + pub async fn reconcile_connected( + &self, + domain: AgentDesiredStateDomain, + ) -> Result<(), AppError> { + match domain { + AgentDesiredStateDomain::PingTasks => { + let _guard = self.locks.ping_tasks.lock().await; + self.reconcile_connected_ping_tasks().await + } + AgentDesiredStateDomain::NetworkProbes => { + let _guard = self.locks.network_probes.lock().await; + self.reconcile_connected_network_probes().await + } + AgentDesiredStateDomain::IpQuality => { + let _guard = self.locks.ip_quality.lock().await; + self.reconcile_connected_ip_quality().await + } + AgentDesiredStateDomain::Firewall => { + let _guard = self.locks.firewall.lock().await; + self.reconcile_connected_firewall().await + } + } + } + + async fn reconcile_connected_ping_tasks(&self) -> Result<(), AppError> { + let server_ids = self.agent_manager.connected_server_ids(); + let tasks = self.load_ping_tasks().await?; + let capabilities = self.capabilities_for_agents(&server_ids).await?; + + for server_id in server_ids { + let server_capabilities = capabilities.get(&server_id).copied().unwrap_or(CAP_DEFAULT); + self.send_ping_tasks(&server_id, &tasks, server_capabilities) + .await; + } + Ok(()) + } + + async fn reconcile_connected_network_probes(&self) -> Result<(), AppError> { + let server_ids = self.agent_manager.connected_server_ids(); + let setting = NetworkProbeService::get_setting(&self.db).await?; + let mut first_error = None; + + for server_id in server_ids { + if let Err(error) = self.send_network_probes(&server_id, &setting).await { + tracing::warn!( + server_id, + error = %error, + "failed to reconcile network probe desired state" + ); + if first_error.is_none() { + first_error = Some(error); + } + } + } + + match first_error { + Some(error) => Err(error), + None => Ok(()), + } + } + + async fn reconcile_connected_ip_quality(&self) -> Result<(), AppError> { + let server_ids = self.agent_manager.connected_server_ids(); + let services = IpQualityService::enabled_service_defs(&self.db).await?; + let interval_hours = ip_quality_interval_hours( + IpQualityService::get_setting(&self.db) + .await? + .check_interval_hours, + )?; + + for server_id in server_ids { + self.send_if_online( + &server_id, + ServerMessage::IpQualitySync { + services: services.clone(), + interval_hours, + }, + ) + .await; + } + Ok(()) + } + + async fn reconcile_connected_firewall(&self) -> Result<(), AppError> { + let server_ids = self.agent_manager.connected_server_ids(); + let mut first_error = None; + + for server_id in server_ids { + if let Err(error) = self.send_firewall(&server_id).await { + tracing::warn!( + server_id, + error = %error, + "failed to reconcile firewall desired state" + ); + if first_error.is_none() { + first_error = Some(error); + } + } + } + + match first_error { + Some(error) => Err(error), + None => Ok(()), + } + } + + async fn load_ping_tasks(&self) -> Result, AppError> { + Ok(ping_task::Entity::find() + .filter(ping_task::Column::Enabled.eq(true)) + .all(&self.db) + .await?) + } + + async fn capabilities_for_agent(&self, server_id: &str) -> Result { + if let Some(capabilities) = self.agent_manager.get_effective_capabilities(server_id) { + return Ok(capabilities); + } + + let mirrored = server::Entity::find_by_id(server_id).one(&self.db).await?; + Ok(mirrored + .and_then(|model| u32::try_from(model.capabilities).ok()) + .unwrap_or(CAP_DEFAULT)) + } + + async fn capabilities_for_agents( + &self, + server_ids: &[String], + ) -> Result, AppError> { + if server_ids.is_empty() { + return Ok(HashMap::new()); + } + + let mirrored: HashMap = server::Entity::find() + .filter(server::Column::Id.is_in(server_ids.iter().cloned())) + .all(&self.db) + .await? + .into_iter() + .filter_map(|model| { + u32::try_from(model.capabilities) + .ok() + .map(|capabilities| (model.id, capabilities)) + }) + .collect(); + + Ok(server_ids + .iter() + .map(|server_id| { + let capabilities = self + .agent_manager + .get_effective_capabilities(server_id) + .or_else(|| mirrored.get(server_id).copied()) + .unwrap_or(CAP_DEFAULT); + (server_id.clone(), capabilities) + }) + .collect()) + } + + async fn agent_has_capability( + &self, + server_id: &str, + capability: u32, + ) -> Result { + Ok(has_capability( + self.capabilities_for_agent(server_id).await?, + capability, + )) + } + + async fn send_ping_tasks( + &self, + server_id: &str, + tasks: &[ping_task::Model], + capabilities: u32, + ) { + let tasks = ping_configs_for_agent(tasks, server_id, capabilities); + self.send_if_online(server_id, ServerMessage::PingTasksSync { tasks }) + .await; + } + + async fn send_network_probes( + &self, + server_id: &str, + setting: &NetworkProbeSetting, + ) -> Result<(), AppError> { + let targets = NetworkProbeService::get_server_targets(&self.db, server_id).await?; + let targets = network_targets_for_agent(targets); + + self.send_if_online( + server_id, + ServerMessage::NetworkProbeSync { + targets, + interval: setting.interval, + packet_count: setting.packet_count, + }, + ) + .await; + Ok(()) + } + + async fn send_ip_quality(&self, server_id: &str) -> Result<(), AppError> { + let services = IpQualityService::enabled_service_defs(&self.db).await?; + let interval_hours = ip_quality_interval_hours( + IpQualityService::get_setting(&self.db) + .await? + .check_interval_hours, + )?; + self.send_if_online( + server_id, + ServerMessage::IpQualitySync { + services, + interval_hours, + }, + ) + .await; + Ok(()) + } + + async fn send_firewall(&self, server_id: &str) -> Result<(), AppError> { + let protocol_version = self + .agent_manager + .get_protocol_version(server_id) + .unwrap_or_default(); + if protocol_version < FIREWALL_MIN_PROTOCOL { + return Ok(()); + } + + // Reset is intentionally sent even when the capability is disabled. + // It removes ServerBee's stale nftables state after capability + // revocation; the agent protocol defines reset as an ungated cleanup. + self.firewall + .push_reset_to(server_id, &self.agent_manager) + .await; + if self + .agent_has_capability(server_id, CAP_FIREWALL_BLOCK) + .await? + { + self.firewall + .push_sync_to(server_id, &self.agent_manager) + .await?; + } + Ok(()) + } + + async fn send_if_online(&self, server_id: &str, message: ServerMessage) { + let Some(sender) = self.agent_manager.get_sender(server_id) else { + return; + }; + if sender.send(message).await.is_err() { + tracing::debug!( + server_id, + "agent disconnected during desired-state reconcile" + ); + } + } +} + +fn ping_configs_for_agent( + tasks: &[ping_task::Model], + server_id: &str, + capabilities: u32, +) -> Vec { + let mut configs = Vec::new(); + + for task in tasks { + let Ok(server_ids) = serde_json::from_str::>(&task.server_ids_json) else { + tracing::warn!( + task_id = task.id, + "ignoring ping task with an invalid persisted server assignment" + ); + continue; + }; + if !server_ids.is_empty() && !server_ids.iter().any(|id| id == server_id) { + continue; + } + + let Some(capability) = probe_type_to_cap(&task.probe_type) else { + tracing::warn!( + task_id = task.id, + probe_type = task.probe_type, + "ignoring ping task with unknown probe type" + ); + continue; + }; + if !has_capability(capabilities, capability) { + continue; + } + + let Ok(interval) = u32::try_from(task.interval) else { + tracing::warn!( + task_id = task.id, + interval = task.interval, + "ignoring ping task with a negative interval" + ); + continue; + }; + configs.push(PingTaskConfig { + task_id: task.id.clone(), + probe_type: task.probe_type.clone(), + target: task.target.clone(), + interval, + }); + } + + configs +} + +fn network_targets_for_agent(targets: Vec) -> Vec { + targets + .into_iter() + .map(|target| NetworkProbeTarget { + target_id: target.id, + name: target.name, + target: target.target, + probe_type: target.probe_type, + }) + .collect() +} + +fn ip_quality_interval_hours(value: i32) -> Result { + u32::try_from(value).map_err(|_| { + tracing::error!( + interval_hours = value, + "invalid persisted IP quality check interval" + ); + AppError::Internal("Invalid persisted IP quality setting".to_string()) + }) +} + +#[cfg(test)] +mod tests { + use std::net::{IpAddr, Ipv4Addr, SocketAddr}; + + use chrono::Utc; + use sea_orm::{ActiveModelTrait, Set}; + use serverbee_common::constants::{ + CAP_FIREWALL_BLOCK, CAP_IP_QUALITY, CAP_PING_HTTP, CAP_PING_ICMP, + }; + use tokio::sync::{broadcast, mpsc}; + + use super::*; + use crate::config::AppConfig; + use crate::test_utils::setup_test_db; + + fn test_addr(port: u16) -> SocketAddr { + SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port) + } + + fn test_reconciler( + db: &DatabaseConnection, + ) -> (AgentDesiredStateReconciler, Arc) { + let (browser_tx, _) = broadcast::channel(16); + let agent_manager = Arc::new(AgentManager::new(browser_tx.clone())); + let firewall = Arc::new(FirewallService::new( + db.clone(), + Arc::new(AppConfig::default()), + browser_tx, + )); + ( + AgentDesiredStateReconciler::new(db.clone(), agent_manager.clone(), firewall), + agent_manager, + ) + } + + async fn receive_messages( + receiver: &mut mpsc::Receiver, + count: usize, + ) -> Vec { + let mut messages = Vec::with_capacity(count); + for _ in 0..count { + messages.push( + tokio::time::timeout(std::time::Duration::from_secs(1), receiver.recv()) + .await + .expect("desired-state message timed out") + .expect("agent channel closed"), + ); + } + messages + } + + #[tokio::test] + async fn connection_reconcile_attempts_every_domain() { + let (db, _tmp) = setup_test_db().await; + let (reconciler, agent_manager) = test_reconciler(&db); + let (sender, mut receiver) = mpsc::channel(16); + agent_manager.add_connection( + "server-1".to_string(), + "Server 1".to_string(), + sender, + test_addr(10001), + ); + agent_manager.update_agent_local_capabilities( + "server-1", + CAP_PING_ICMP | CAP_IP_QUALITY | CAP_FIREWALL_BLOCK, + ); + agent_manager.set_protocol_version("server-1", FIREWALL_MIN_PROTOCOL); + + reconciler.reconcile_connection("server-1").await.unwrap(); + + let messages = receive_messages(&mut receiver, 5).await; + assert!(messages.iter().any( + |message| matches!(message, ServerMessage::PingTasksSync { tasks } if tasks.is_empty()) + )); + assert!(messages.iter().any( + |message| matches!(message, ServerMessage::NetworkProbeSync { targets, .. } if targets.is_empty()) + )); + assert!( + messages + .iter() + .any(|message| matches!(message, ServerMessage::IpQualitySync { .. })) + ); + assert!( + messages + .iter() + .any(|message| matches!(message, ServerMessage::BlocklistReset)) + ); + assert!(messages.iter().any( + |message| matches!(message, ServerMessage::BlocklistSync { entries } if entries.is_empty()) + )); + } + + #[tokio::test] + async fn ping_reconcile_filters_scope_and_live_capabilities() { + let (db, _tmp) = setup_test_db().await; + ping_task::ActiveModel { + id: Set("http-task".to_string()), + name: Set("HTTP".to_string()), + probe_type: Set("http".to_string()), + target: Set("https://example.com".to_string()), + interval: Set(60), + server_ids_json: Set("[]".to_string()), + enabled: Set(true), + created_at: Set(Utc::now()), + } + .insert(&db) + .await + .unwrap(); + ping_task::ActiveModel { + id: Set("icmp-task".to_string()), + name: Set("ICMP".to_string()), + probe_type: Set("icmp".to_string()), + target: Set("1.1.1.1".to_string()), + interval: Set(60), + server_ids_json: Set("[\"another-server\"]".to_string()), + enabled: Set(true), + created_at: Set(Utc::now()), + } + .insert(&db) + .await + .unwrap(); + ping_task::ActiveModel { + id: Set("corrupt-task".to_string()), + name: Set("Corrupt".to_string()), + probe_type: Set("http".to_string()), + target: Set("https://example.net".to_string()), + interval: Set(60), + server_ids_json: Set("not-json".to_string()), + enabled: Set(true), + created_at: Set(Utc::now()), + } + .insert(&db) + .await + .unwrap(); + + let (reconciler, agent_manager) = test_reconciler(&db); + let (sender, mut receiver) = mpsc::channel(4); + agent_manager.add_connection( + "server-1".to_string(), + "Server 1".to_string(), + sender, + test_addr(10002), + ); + agent_manager.update_agent_local_capabilities("server-1", CAP_PING_HTTP); + + reconciler + .reconcile_agent("server-1", AgentDesiredStateDomain::PingTasks) + .await + .unwrap(); + + let message = receiver.recv().await.unwrap(); + let ServerMessage::PingTasksSync { tasks } = message else { + panic!("expected ping task sync"); + }; + assert_eq!(tasks.len(), 1); + assert_eq!(tasks[0].task_id, "http-task"); + } + + #[tokio::test] + async fn connected_reconcile_only_sends_the_requested_domain() { + let (db, _tmp) = setup_test_db().await; + let (reconciler, agent_manager) = test_reconciler(&db); + let (sender_a, mut receiver_a) = mpsc::channel(4); + let (sender_b, mut receiver_b) = mpsc::channel(4); + agent_manager.add_connection( + "server-a".to_string(), + "Server A".to_string(), + sender_a, + test_addr(10003), + ); + agent_manager.add_connection( + "server-b".to_string(), + "Server B".to_string(), + sender_b, + test_addr(10004), + ); + agent_manager.update_agent_local_capabilities("server-a", CAP_PING_ICMP); + agent_manager.update_agent_local_capabilities("server-b", CAP_PING_ICMP); + + reconciler + .reconcile_connected(AgentDesiredStateDomain::NetworkProbes) + .await + .unwrap(); + + assert!(matches!( + receiver_a.recv().await.unwrap(), + ServerMessage::NetworkProbeSync { .. } + )); + assert!(matches!( + receiver_b.recv().await.unwrap(), + ServerMessage::NetworkProbeSync { .. } + )); + assert!(receiver_a.try_recv().is_err()); + assert!(receiver_b.try_recv().is_err()); + } + + #[tokio::test] + async fn firewall_reconcile_resets_state_after_capability_revocation() { + let (db, _tmp) = setup_test_db().await; + let (reconciler, agent_manager) = test_reconciler(&db); + let (sender, mut receiver) = mpsc::channel(4); + agent_manager.add_connection( + "server-1".to_string(), + "Server 1".to_string(), + sender, + test_addr(10005), + ); + agent_manager.update_agent_local_capabilities("server-1", 0); + agent_manager.set_protocol_version("server-1", FIREWALL_MIN_PROTOCOL); + + reconciler + .reconcile_agent("server-1", AgentDesiredStateDomain::Firewall) + .await + .unwrap(); + + assert!(matches!( + receiver.recv().await.unwrap(), + ServerMessage::BlocklistReset + )); + assert!(receiver.try_recv().is_err()); + } +} diff --git a/crates/server/src/service/firewall.rs b/crates/server/src/service/firewall.rs index 7ea852039..6b445078d 100644 --- a/crates/server/src/service/firewall.rs +++ b/crates/server/src/service/firewall.rs @@ -305,9 +305,12 @@ impl FirewallService { Ok(()) } - /// Tell an agent to drop every entry it currently holds and clear our - /// `apply_state` map for that agent. Caller is responsible for verifying - /// capability + protocol version. + /// Tell an agent to drop every entry it currently holds. Caller is + /// responsible for verifying capability + protocol version. + /// + /// Sending the command is not proof that the kernel state was wiped, so + /// `apply_state` remains untouched until [`Self::record_reset_ack`] receives + /// a successful acknowledgement. pub async fn push_reset_to( &self, server_id: &str, @@ -318,8 +321,6 @@ impl FirewallService { .send(serverbee_common::protocol::ServerMessage::BlocklistReset) .await; } - let mut g = self.apply_state.write().await; - g.retain(|(_block_id, srv), _| srv != server_id); } /// Update `apply_state`, write an audit row, and fan out a @@ -510,8 +511,10 @@ impl FirewallService { Ok(Some(row.id)) } - /// Audit a `BlocklistResetAck` from the agent. No `apply_state` mutation - /// is needed because `push_reset_to` already cleared it locally. + /// Apply and audit a `BlocklistResetAck` from the agent. A successful ack is + /// the only evidence that the agent wiped its kernel state, so only that + /// path clears the server's entries from `apply_state`. Failed acks retain + /// the last-known state for operators and later retries. pub async fn record_reset_ack( &self, server_id: &str, @@ -519,6 +522,11 @@ impl FirewallService { reason: Option, db: &DatabaseConnection, ) { + if ok { + let mut g = self.apply_state.write().await; + g.retain(|(_block_id, srv), _| srv != server_id); + } + let action = if ok { "firewall_reset_acked" } else { @@ -1219,7 +1227,7 @@ mod tests { } #[tokio::test] - async fn push_reset_sends_reset_and_clears_only_target_apply_state() { + async fn push_reset_sends_reset_without_clearing_apply_state() { use serverbee_common::firewall::FIREWALL_MIN_PROTOCOL; use serverbee_common::protocol::ServerMessage; @@ -1227,7 +1235,8 @@ mod tests { let (svc, mgr) = make_service(db); let mut rx = wire_agent(&mgr, "srv-A", FIREWALL_MIN_PROTOCOL, 0); - // Seed apply_state for srv-A and srv-B; only srv-A should be cleared. + // Sending the command is not proof that the agent wiped its state. + // Both entries must remain until a successful reset acknowledgement. { let mut g = svc.apply_state.write().await; g.insert( @@ -1254,12 +1263,12 @@ mod tests { assert!(matches!(msg, ServerMessage::BlocklistReset)); let g = svc.apply_state.read().await; - assert!(!g.contains_key(&("blk-1".into(), "srv-A".to_string()))); + assert!(g.contains_key(&("blk-1".into(), "srv-A".to_string()))); assert!(g.contains_key(&("blk-2".into(), "srv-B".to_string()))); } #[tokio::test] - async fn push_reset_clears_state_even_without_connection() { + async fn push_reset_without_connection_preserves_apply_state() { let (db, _tmp) = setup_test_db().await; let (svc, mgr) = make_service(db); { @@ -1273,10 +1282,11 @@ mod tests { }, ); } - // No connection for "ghost" — the local apply_state is still cleared. + // With no connection there can be no successful acknowledgement, so + // the last-known state must remain intact. svc.push_reset_to("ghost", &mgr).await; let g = svc.apply_state.read().await; - assert!(g.is_empty()); + assert!(g.contains_key(&("blk-1".into(), "ghost".to_string()))); } // ── record_ack: state mirror + audit + broadcast ── @@ -1366,25 +1376,68 @@ mod tests { // ── record_reset_ack: ok / failure audit actions ── #[tokio::test] - async fn record_reset_ack_ok_writes_acked_audit() { + async fn record_reset_ack_ok_writes_acked_audit_and_clears_only_target_state() { let (db, _tmp) = setup_test_db().await; let (svc, _mgr) = make_service(db.clone()); + + { + let mut g = svc.apply_state.write().await; + g.insert( + ("blk-1".into(), "srv-A".into()), + ApplyState { + state: BlocklistEntryState::Present, + reason: None, + at: chrono::Utc::now(), + }, + ); + g.insert( + ("blk-2".into(), "srv-B".into()), + ApplyState { + state: BlocklistEntryState::Present, + reason: None, + at: chrono::Utc::now(), + }, + ); + } + svc.record_reset_ack("srv-A", true, None, &db).await; + let acked = fetch_audits(&db, "firewall_reset_acked").await; assert_eq!(acked.len(), 1); + + let g = svc.apply_state.read().await; + assert!(!g.contains_key(&("blk-1".into(), "srv-A".to_string()))); + assert!(g.contains_key(&("blk-2".into(), "srv-B".to_string()))); } #[tokio::test] - async fn record_reset_ack_failure_writes_failed_audit() { + async fn record_reset_ack_failure_writes_failed_audit_and_preserves_state() { let (db, _tmp) = setup_test_db().await; let (svc, _mgr) = make_service(db.clone()); + + { + let mut g = svc.apply_state.write().await; + g.insert( + ("blk-1".into(), "srv-A".into()), + ApplyState { + state: BlocklistEntryState::Present, + reason: None, + at: chrono::Utc::now(), + }, + ); + } + svc.record_reset_ack("srv-A", false, Some("agent error".into()), &db) .await; + let failed = fetch_audits(&db, "firewall_reset_failed_agent").await; assert_eq!(failed.len(), 1); // The success action must NOT be written. let acked = fetch_audits(&db, "firewall_reset_acked").await; assert!(acked.is_empty()); + + let g = svc.apply_state.read().await; + assert!(g.contains_key(&("blk-1".into(), "srv-A".to_string()))); } // ── broadcast_changed_* ── diff --git a/crates/server/src/service/mod.rs b/crates/server/src/service/mod.rs index 0c10b63a8..b24fef922 100644 --- a/crates/server/src/service/mod.rs +++ b/crates/server/src/service/mod.rs @@ -1,4 +1,5 @@ pub mod agent_manager; +pub mod agent_reconcile; pub mod alert; pub mod apns; pub mod asn; diff --git a/crates/server/src/service/ping.rs b/crates/server/src/service/ping.rs index d7b0ff8c1..5229bf0a0 100644 --- a/crates/server/src/service/ping.rs +++ b/crates/server/src/service/ping.rs @@ -3,12 +3,8 @@ use sea_orm::*; use serde::{Deserialize, Serialize}; use uuid::Uuid; -use crate::entity::{ping_record, ping_task, server}; +use crate::entity::{ping_record, ping_task}; use crate::error::AppError; -use crate::service::agent_manager::AgentManager; -use serverbee_common::constants::{CAP_DEFAULT, has_capability, probe_type_to_cap}; -use serverbee_common::protocol::ServerMessage; -use serverbee_common::types::PingTaskConfig; pub struct PingService; @@ -57,7 +53,6 @@ impl PingService { pub async fn create( db: &DatabaseConnection, - agent_manager: &AgentManager, input: CreatePingTask, ) -> Result { if !["icmp", "tcp", "http"].contains(&input.probe_type.as_str()) { @@ -81,17 +76,11 @@ impl PingService { enabled: Set(input.enabled), created_at: Set(Utc::now()), }; - let created = model.insert(db).await?; - - // Sync tasks to affected agents - Self::sync_tasks_to_agents(db, agent_manager).await; - - Ok(created) + Ok(model.insert(db).await?) } pub async fn update( db: &DatabaseConnection, - agent_manager: &AgentManager, id: &str, input: UpdatePingTask, ) -> Result { @@ -126,18 +115,10 @@ impl PingService { model.enabled = Set(enabled); } - let updated = model.update(db).await?; - - Self::sync_tasks_to_agents(db, agent_manager).await; - - Ok(updated) + Ok(model.update(db).await?) } - pub async fn delete( - db: &DatabaseConnection, - agent_manager: &AgentManager, - id: &str, - ) -> Result<(), AppError> { + pub async fn delete(db: &DatabaseConnection, id: &str) -> Result<(), AppError> { let result = ping_task::Entity::delete_by_id(id).exec(db).await?; if result.rows_affected == 0 { return Err(AppError::NotFound(format!("Ping task {id} not found"))); @@ -148,8 +129,6 @@ impl PingService { .exec(db) .await?; - Self::sync_tasks_to_agents(db, agent_manager).await; - Ok(()) } @@ -174,155 +153,12 @@ impl PingService { .all(db) .await?) } - - /// Send current ping tasks to a specific agent (e.g., on new connection). - pub async fn sync_tasks_to_agent( - db: &DatabaseConnection, - agent_manager: &AgentManager, - server_id: &str, - ) { - let tasks = match ping_task::Entity::find() - .filter(ping_task::Column::Enabled.eq(true)) - .all(db) - .await - { - Ok(t) => t, - Err(e) => { - tracing::error!("Failed to load ping tasks for agent sync: {e}"); - return; - } - }; - - // Fetch server capabilities - let server_caps = server::Entity::find_by_id(server_id) - .one(db) - .await - .ok() - .flatten() - .map(|s| s.capabilities as u32) - .unwrap_or(CAP_DEFAULT); - - let mut task_configs: Vec = Vec::new(); - for task in &tasks { - let server_ids: Vec = - serde_json::from_str(&task.server_ids_json).unwrap_or_default(); - // Include task if server_ids is empty (all agents) or contains this server - if server_ids.is_empty() || server_ids.contains(&server_id.to_string()) { - // Filter by capability - if probe_type_to_cap(&task.probe_type) - .map(|cap| has_capability(server_caps, cap)) - .unwrap_or(false) - { - task_configs.push(PingTaskConfig { - task_id: task.id.clone(), - probe_type: task.probe_type.clone(), - target: task.target.clone(), - interval: task.interval as u32, - }); - } - } - } - - // Always send PingTasksSync (even if empty — tells Agent to stop all probes) - if let Some(tx) = agent_manager.get_sender(server_id) { - let msg = ServerMessage::PingTasksSync { - tasks: task_configs, - }; - let _ = tx.send(msg).await; - } - } - - /// Sync all enabled ping tasks to all connected agents. - async fn sync_tasks_to_agents(db: &DatabaseConnection, agent_manager: &AgentManager) { - let tasks = match ping_task::Entity::find() - .filter(ping_task::Column::Enabled.eq(true)) - .all(db) - .await - { - Ok(t) => t, - Err(e) => { - tracing::error!("Failed to load ping tasks for sync: {e}"); - return; - } - }; - - // Fetch capabilities for all connected agents - let connected_ids = agent_manager.connected_server_ids(); - let server_caps_map: std::collections::HashMap = match server::Entity::find() - .filter(server::Column::Id.is_in(connected_ids.iter().cloned())) - .all(db) - .await - { - Ok(servers) => servers - .into_iter() - .map(|s| { - let caps = s.capabilities as u32; - (s.id, caps) - }) - .collect(), - Err(e) => { - tracing::error!("Failed to load server caps for ping sync: {e}"); - return; - } - }; - - // Build per-agent task lists filtered by capability - let mut agent_tasks: std::collections::HashMap> = - std::collections::HashMap::new(); - - // Ensure every connected agent gets an entry (even if empty) - for sid in &connected_ids { - agent_tasks.entry(sid.clone()).or_default(); - } - - for task in &tasks { - let server_ids: Vec = - serde_json::from_str(&task.server_ids_json).unwrap_or_default(); - let config = PingTaskConfig { - task_id: task.id.clone(), - probe_type: task.probe_type.clone(), - target: task.target.clone(), - interval: task.interval as u32, - }; - - let target_ids: Vec = if server_ids.is_empty() { - connected_ids.clone() - } else { - server_ids - }; - - for sid in target_ids { - let caps = server_caps_map.get(&sid).copied().unwrap_or(CAP_DEFAULT); - if probe_type_to_cap(&task.probe_type) - .map(|cap| has_capability(caps, cap)) - .unwrap_or(false) - { - agent_tasks.entry(sid).or_default().push(config.clone()); - } - } - } - - for (server_id, task_configs) in agent_tasks { - if let Some(tx) = agent_manager.get_sender(&server_id) { - let msg = ServerMessage::PingTasksSync { - tasks: task_configs, - }; - let _ = tx.send(msg).await; - } - } - } } #[cfg(test)] mod tests { use super::*; use crate::test_utils::setup_test_db; - use serverbee_common::constants::{CAP_PING_HTTP, CAP_PING_ICMP}; - - fn test_agent_manager() -> crate::service::agent_manager::AgentManager { - let (tx, _) = tokio::sync::broadcast::channel(16); - crate::service::agent_manager::AgentManager::new(tx) - } fn sample_create_ping_task() -> CreatePingTask { CreatePingTask { @@ -338,12 +174,9 @@ mod tests { #[tokio::test] async fn test_create_and_list_ping_task() { let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); let input = sample_create_ping_task(); - let created = PingService::create(&db, &agent_manager, input) - .await - .unwrap(); + let created = PingService::create(&db, input).await.unwrap(); let list = PingService::list(&db).await.unwrap(); assert_eq!(list.len(), 1); @@ -355,13 +188,12 @@ mod tests { #[tokio::test] async fn create_rejects_literal_metadata_target() { let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); let mut input = sample_create_ping_task(); input.probe_type = "http".to_string(); input.target = "http://169.254.169.254/latest/meta-data/".to_string(); - let result = PingService::create(&db, &agent_manager, input).await; + let result = PingService::create(&db, input).await; assert!( result.is_err(), "creating a ping task targeting cloud metadata must be rejected" @@ -372,15 +204,13 @@ mod tests { #[tokio::test] async fn update_rejects_literal_loopback_target() { let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); - let created = PingService::create(&db, &agent_manager, sample_create_ping_task()) + let created = PingService::create(&db, sample_create_ping_task()) .await .unwrap(); let result = PingService::update( &db, - &agent_manager, &created.id, UpdatePingTask { name: None, @@ -401,16 +231,11 @@ mod tests { #[tokio::test] async fn test_delete_ping_task() { let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); let input = sample_create_ping_task(); - let created = PingService::create(&db, &agent_manager, input) - .await - .unwrap(); + let created = PingService::create(&db, input).await.unwrap(); - PingService::delete(&db, &agent_manager, &created.id) - .await - .unwrap(); + PingService::delete(&db, &created.id).await.unwrap(); let list = PingService::list(&db).await.unwrap(); assert!(list.is_empty()); @@ -419,12 +244,9 @@ mod tests { #[tokio::test] async fn test_get_ping_task() { let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); let input = sample_create_ping_task(); - let created = PingService::create(&db, &agent_manager, input) - .await - .unwrap(); + let created = PingService::create(&db, input).await.unwrap(); let fetched = PingService::get(&db, &created.id).await.unwrap(); assert_eq!(fetched.id, created.id); @@ -435,25 +257,6 @@ mod tests { // --- helpers -------------------------------------------------------- - /// Seed a minimal `servers` row with the given capabilities mirror. - async fn insert_test_server(db: &DatabaseConnection, id: &str, capabilities: u32) { - let now = Utc::now(); - server::ActiveModel { - id: Set(id.to_string()), - name: Set(format!("srv-{id}")), - weight: Set(0), - hidden: Set(false), - capabilities: Set(capabilities as i32), - protocol_version: Set(1), - created_at: Set(now), - updated_at: Set(now), - ..Default::default() - } - .insert(db) - .await - .expect("insert test server should succeed"); - } - /// Seed a single ping record and return the inserted row id. async fn insert_ping_record( db: &DatabaseConnection, @@ -489,19 +292,6 @@ mod tests { } } - /// Register a connected agent with a real mpsc sender so `get_sender` - /// returns `Some`. Returns the receiver so the channel stays open and we - /// can inspect delivered `ServerMessage`s. - fn connect_agent( - agent_manager: &crate::service::agent_manager::AgentManager, - server_id: &str, - ) -> tokio::sync::mpsc::Receiver { - let (tx, rx) = tokio::sync::mpsc::channel(16); - let addr = std::net::SocketAddr::from(([127, 0, 0, 1], 8080)); - agent_manager.add_connection(server_id.to_string(), format!("srv-{server_id}"), tx, addr); - rx - } - // --- get error path ------------------------------------------------- #[tokio::test] @@ -516,12 +306,11 @@ mod tests { #[tokio::test] async fn create_rejects_invalid_probe_type() { let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); let mut input = sample_create_ping_task(); input.probe_type = "udp".to_string(); - let result = PingService::create(&db, &agent_manager, input).await; + let result = PingService::create(&db, input).await; assert!(matches!(result, Err(AppError::Validation(_)))); // Nothing should have been persisted. assert!(PingService::list(&db).await.unwrap().is_empty()); @@ -530,16 +319,13 @@ mod tests { #[tokio::test] async fn create_persists_server_ids_json() { let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); let mut input = sample_create_ping_task(); input.probe_type = "icmp".to_string(); input.target = "1.1.1.1".to_string(); input.server_ids = vec!["s1".to_string(), "s2".to_string()]; - let created = PingService::create(&db, &agent_manager, input) - .await - .unwrap(); + let created = PingService::create(&db, input).await.unwrap(); let parsed: Vec = serde_json::from_str(&created.server_ids_json).unwrap(); assert_eq!(parsed, vec!["s1".to_string(), "s2".to_string()]); } @@ -549,24 +335,21 @@ mod tests { #[tokio::test] async fn update_missing_task_returns_not_found() { let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); - let result = PingService::update(&db, &agent_manager, "ghost", empty_update()).await; + let result = PingService::update(&db, "ghost", empty_update()).await; assert!(matches!(result, Err(AppError::NotFound(_)))); } #[tokio::test] async fn update_applies_all_fields() { let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); - let created = PingService::create(&db, &agent_manager, sample_create_ping_task()) + let created = PingService::create(&db, sample_create_ping_task()) .await .unwrap(); let updated = PingService::update( &db, - &agent_manager, &created.id, UpdatePingTask { name: Some("Renamed".to_string()), @@ -592,13 +375,12 @@ mod tests { #[tokio::test] async fn update_with_no_fields_is_noop() { let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); - let created = PingService::create(&db, &agent_manager, sample_create_ping_task()) + let created = PingService::create(&db, sample_create_ping_task()) .await .unwrap(); - let updated = PingService::update(&db, &agent_manager, &created.id, empty_update()) + let updated = PingService::update(&db, &created.id, empty_update()) .await .unwrap(); @@ -613,15 +395,13 @@ mod tests { #[tokio::test] async fn update_rejects_invalid_probe_type() { let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); - let created = PingService::create(&db, &agent_manager, sample_create_ping_task()) + let created = PingService::create(&db, sample_create_ping_task()) .await .unwrap(); let result = PingService::update( &db, - &agent_manager, &created.id, UpdatePingTask { probe_type: Some("ftp".to_string()), @@ -641,18 +421,16 @@ mod tests { #[tokio::test] async fn delete_missing_task_returns_not_found() { let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); - let result = PingService::delete(&db, &agent_manager, "ghost").await; + let result = PingService::delete(&db, "ghost").await; assert!(matches!(result, Err(AppError::NotFound(_)))); } #[tokio::test] async fn delete_removes_associated_records() { let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); - let created = PingService::create(&db, &agent_manager, sample_create_ping_task()) + let created = PingService::create(&db, sample_create_ping_task()) .await .unwrap(); @@ -662,9 +440,7 @@ mod tests { // A record for an unrelated task must survive the delete. insert_ping_record(&db, "other-task", "s1", 99.0, true, t).await; - PingService::delete(&db, &agent_manager, &created.id) - .await - .unwrap(); + PingService::delete(&db, &created.id).await.unwrap(); let remaining = ping_record::Entity::find().all(&db).await.unwrap(); assert_eq!(remaining.len(), 1); @@ -745,263 +521,4 @@ mod tests { .unwrap(); assert!(records.is_empty()); } - - // --- sync_tasks_to_agent (single) ----------------------------------- - - /// Drain all currently-buffered messages from the agent's receiver. - fn drain_messages(rx: &mut tokio::sync::mpsc::Receiver) -> Vec { - let mut out = Vec::new(); - while let Ok(msg) = rx.try_recv() { - out.push(msg); - } - out - } - - fn ping_tasks_from(msgs: &[ServerMessage]) -> Option> { - msgs.iter().find_map(|m| match m { - ServerMessage::PingTasksSync { tasks } => Some(tasks.clone()), - _ => None, - }) - } - - #[tokio::test] - async fn sync_to_agent_includes_capable_tasks_and_skips_uncapable() { - let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); - - // Server supports only ICMP ping (and nothing else relevant). - insert_test_server(&db, "s1", CAP_PING_ICMP).await; - - // ICMP task targeting all agents (empty server_ids) -> should be sent. - PingService::create( - &db, - &agent_manager, - CreatePingTask { - name: "icmp-all".to_string(), - probe_type: "icmp".to_string(), - target: "1.1.1.1".to_string(), - interval: 30, - server_ids: vec![], - enabled: true, - }, - ) - .await - .unwrap(); - // HTTP task -> server lacks CAP_PING_HTTP, must be filtered out. - PingService::create( - &db, - &agent_manager, - CreatePingTask { - name: "http-all".to_string(), - probe_type: "http".to_string(), - target: "https://example.com".to_string(), - interval: 30, - server_ids: vec![], - enabled: true, - }, - ) - .await - .unwrap(); - - let mut rx = connect_agent(&agent_manager, "s1"); - PingService::sync_tasks_to_agent(&db, &agent_manager, "s1").await; - - let msgs = drain_messages(&mut rx); - let tasks = ping_tasks_from(&msgs).expect("a PingTasksSync must be delivered"); - assert_eq!(tasks.len(), 1, "only the ICMP task is capability-allowed"); - assert_eq!(tasks[0].probe_type, "icmp"); - assert_eq!(tasks[0].interval, 30); - } - - #[tokio::test] - async fn sync_to_agent_respects_server_id_targeting() { - let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); - - insert_test_server(&db, "s1", CAP_DEFAULT).await; - - // Task scoped to a different server -> not for s1. - PingService::create( - &db, - &agent_manager, - CreatePingTask { - name: "scoped-other".to_string(), - probe_type: "icmp".to_string(), - target: "1.1.1.1".to_string(), - interval: 60, - server_ids: vec!["s2".to_string()], - enabled: true, - }, - ) - .await - .unwrap(); - // Task explicitly scoped to s1 -> included. - PingService::create( - &db, - &agent_manager, - CreatePingTask { - name: "scoped-s1".to_string(), - probe_type: "tcp".to_string(), - target: "1.1.1.1:53".to_string(), - interval: 60, - server_ids: vec!["s1".to_string()], - enabled: true, - }, - ) - .await - .unwrap(); - - let mut rx = connect_agent(&agent_manager, "s1"); - PingService::sync_tasks_to_agent(&db, &agent_manager, "s1").await; - - let msgs = drain_messages(&mut rx); - let tasks = ping_tasks_from(&msgs).expect("a PingTasksSync must be delivered"); - assert_eq!(tasks.len(), 1); - assert_eq!(tasks[0].probe_type, "tcp"); - } - - #[tokio::test] - async fn sync_to_agent_excludes_disabled_tasks() { - let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); - - insert_test_server(&db, "s1", CAP_DEFAULT).await; - - PingService::create( - &db, - &agent_manager, - CreatePingTask { - name: "disabled".to_string(), - probe_type: "icmp".to_string(), - target: "1.1.1.1".to_string(), - interval: 60, - server_ids: vec![], - enabled: false, - }, - ) - .await - .unwrap(); - - let mut rx = connect_agent(&agent_manager, "s1"); - PingService::sync_tasks_to_agent(&db, &agent_manager, "s1").await; - - let msgs = drain_messages(&mut rx); - // Sync still fires (telling the agent to stop probes) but with no tasks. - let tasks = ping_tasks_from(&msgs).expect("a PingTasksSync must be delivered"); - assert!(tasks.is_empty()); - } - - #[tokio::test] - async fn sync_to_agent_uses_cap_default_when_server_row_missing() { - let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); - - // No `servers` row for s1 -> falls back to CAP_DEFAULT (ICMP allowed). - PingService::create( - &db, - &agent_manager, - CreatePingTask { - name: "icmp".to_string(), - probe_type: "icmp".to_string(), - target: "1.1.1.1".to_string(), - interval: 60, - server_ids: vec![], - enabled: true, - }, - ) - .await - .unwrap(); - - let mut rx = connect_agent(&agent_manager, "s1"); - PingService::sync_tasks_to_agent(&db, &agent_manager, "s1").await; - - let msgs = drain_messages(&mut rx); - let tasks = ping_tasks_from(&msgs).expect("a PingTasksSync must be delivered"); - assert_eq!(tasks.len(), 1); - assert_eq!(tasks[0].probe_type, "icmp"); - } - - #[tokio::test] - async fn sync_to_agent_noop_when_agent_not_connected() { - let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); - - insert_test_server(&db, "s1", CAP_DEFAULT).await; - PingService::create(&db, &agent_manager, sample_create_ping_task()) - .await - .unwrap(); - - // No connection registered -> get_sender returns None, nothing to assert - // beyond "does not panic". - PingService::sync_tasks_to_agent(&db, &agent_manager, "s1").await; - } - - // --- sync_tasks_to_agents (plural, exercised via create/update/delete) - - - #[tokio::test] - async fn create_syncs_per_agent_filtered_by_capability() { - let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); - - // s1 supports HTTP ping, s2 does not. - insert_test_server(&db, "s1", CAP_PING_HTTP).await; - insert_test_server(&db, "s2", CAP_PING_ICMP).await; - let mut rx1 = connect_agent(&agent_manager, "s1"); - let mut rx2 = connect_agent(&agent_manager, "s2"); - - // Creating the task triggers sync_tasks_to_agents over all connected agents. - PingService::create( - &db, - &agent_manager, - CreatePingTask { - name: "http-all".to_string(), - probe_type: "http".to_string(), - target: "https://example.com".to_string(), - interval: 45, - server_ids: vec![], - enabled: true, - }, - ) - .await - .unwrap(); - - let tasks1 = ping_tasks_from(&drain_messages(&mut rx1)).expect("s1 gets a sync"); - assert_eq!(tasks1.len(), 1, "s1 has HTTP capability"); - assert_eq!(tasks1[0].interval, 45); - - let tasks2 = ping_tasks_from(&drain_messages(&mut rx2)).expect("s2 gets a sync"); - assert!(tasks2.is_empty(), "s2 lacks HTTP capability"); - } - - #[tokio::test] - async fn create_syncs_explicit_server_ids_to_targeted_agent_only() { - let (db, _tmp) = setup_test_db().await; - let agent_manager = test_agent_manager(); - - insert_test_server(&db, "s1", CAP_DEFAULT).await; - insert_test_server(&db, "s2", CAP_DEFAULT).await; - let mut rx1 = connect_agent(&agent_manager, "s1"); - let mut rx2 = connect_agent(&agent_manager, "s2"); - - PingService::create( - &db, - &agent_manager, - CreatePingTask { - name: "scoped-s1".to_string(), - probe_type: "icmp".to_string(), - target: "1.1.1.1".to_string(), - interval: 60, - server_ids: vec!["s1".to_string()], - enabled: true, - }, - ) - .await - .unwrap(); - - let tasks1 = ping_tasks_from(&drain_messages(&mut rx1)).expect("s1 gets a sync"); - assert_eq!(tasks1.len(), 1); - - let tasks2 = ping_tasks_from(&drain_messages(&mut rx2)).expect("s2 gets a sync"); - assert!(tasks2.is_empty(), "task is scoped to s1 only"); - } } diff --git a/crates/server/src/state.rs b/crates/server/src/state.rs index 6e837a6df..c568a4551 100644 --- a/crates/server/src/state.rs +++ b/crates/server/src/state.rs @@ -9,16 +9,17 @@ use tokio::sync::broadcast; use crate::config::AppConfig; use crate::error::AppError; use crate::service::agent_manager::AgentManager; +use crate::service::agent_reconcile::AgentDesiredStateReconciler; use crate::service::alert::AlertStateManager; use crate::service::asn::AsnService; use crate::service::docker_viewer::DockerViewerTracker; use crate::service::file_transfer::FileTransferManager; +use crate::service::firewall::FirewallService; use crate::service::geoip::GeoIpService; use crate::service::high_risk_audit::{ DockerLogsAuditContext, ExecAuditContext, TerminalAuditContext, }; use crate::service::monitor_check::MonitorCheckRunner; -use crate::service::firewall::FirewallService; use crate::service::security::SecurityService; use crate::service::task_scheduler::TaskScheduler; use crate::service::upgrade_release::UpgradeReleaseService; @@ -105,6 +106,8 @@ pub struct AppState { /// Firewall blocklist service. CRUD wiring lives in /// `router::api::firewall`; WS push is invoked from there. pub firewall: Arc, + /// Projects server-owned desired state into full-state agent messages. + pub agent_desired_state: AgentDesiredStateReconciler, /// Pending mobile pairing codes for QR login, keyed by code. pub pending_pairs: DashMap, /// Terminal session audit contexts keyed by session_id. @@ -231,6 +234,8 @@ impl AppState { config_arc.clone(), browser_tx.clone(), )); + let agent_desired_state = + AgentDesiredStateReconciler::new(db.clone(), agent_manager.clone(), firewall.clone()); let security_service = Arc::new(SecurityService::new( db.clone(), browser_tx.clone(), @@ -264,6 +269,7 @@ impl AppState { alert_state_manager, security_service, firewall, + agent_desired_state, pending_pairs: DashMap::new(), terminal_audit_contexts: DashMap::new(), docker_logs_audit_contexts: DashMap::new(), diff --git a/crates/server/tests/agent_messages.rs b/crates/server/tests/agent_messages.rs index 9253d2414..c3377e307 100644 --- a/crates/server/tests/agent_messages.rs +++ b/crates/server/tests/agent_messages.rs @@ -25,12 +25,14 @@ mod common; use common::{ - connect_agent, http_client, login_admin, recv_agent_text, register_agent, send_system_info, - start_test_server, AgentReader, AgentSink, + AgentReader, AgentSink, connect_agent, http_client, login_admin, recv_agent_text, + register_agent, send_system_info, start_test_server, }; use futures_util::SinkExt; -use serde_json::{json, Value}; -use serverbee_common::constants::CAP_DEFAULT; +use serde_json::{Value, json}; +use serverbee_common::constants::{ + CAP_DEFAULT, CAP_FIREWALL_BLOCK, CAP_PING_HTTP, CAP_PING_ICMP, CAP_PING_TCP, +}; use tokio_tungstenite::tungstenite; // --------------------------------------------------------------------------- @@ -659,7 +661,101 @@ async fn test_capabilities_changed_is_audited_and_mirrored() { ); // Connection survives. - send_system_info(&mut sink, &mut reader, "post-capchange-handshake", Some(new_caps)).await; + send_system_info( + &mut sink, + &mut reader, + "post-capchange-handshake", + Some(new_caps), + ) + .await; + + let _ = sink.close().await; +} + +#[tokio::test] +async fn test_capabilities_changed_reconciles_agent_desired_state_after_revoke() { + let (base_url, _tmp) = start_test_server().await; + let client = http_client(); + login_admin(&client, &base_url).await; + + let create_resp = client + .post(format!("{}/api/ping-tasks", base_url)) + .json(&json!({ + "name": "capability-reconcile-icmp", + "probe_type": "icmp", + "target": "1.1.1.1", + "interval": 60, + "server_ids": [] + })) + .send() + .await + .expect("create ping task failed"); + assert_eq!( + create_resp.status(), + 200, + "ping task fixture should be created" + ); + + let (_server_id, mut sink, mut reader) = bring_up_agent(&client, &base_url).await; + + let revoked_caps = + CAP_DEFAULT & !(CAP_PING_ICMP | CAP_PING_TCP | CAP_PING_HTTP | CAP_FIREWALL_BLOCK); + send_agent_frame( + &mut sink, + json!({ + "type": "capabilities_changed", + "msg_id": "cap-reconcile-revoke-1", + "capabilities": revoked_caps, + "temporary": [], + "changes": [{ + "cap": "firewall_block", + "action": "revoked", + "reason": "policy_changed" + }] + }), + ) + .await; + + let expected = [ + "ping_tasks_sync", + "network_probe_sync", + "ip_quality_sync", + "blocklist_reset", + ]; + let mut seen = std::collections::HashSet::new(); + let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5); + + while seen.len() < expected.len() { + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + assert!( + !remaining.is_zero(), + "timed out waiting for capability-revoke reconciliation; saw {seen:?}" + ); + let message = tokio::time::timeout(remaining, recv_agent_text(&mut reader)) + .await + .unwrap_or_else(|_| { + panic!("timed out waiting for capability-revoke reconciliation; saw {seen:?}") + }); + let Some(message_type) = message["type"].as_str() else { + continue; + }; + if expected.contains(&message_type) { + if message_type == "ping_tasks_sync" { + assert_eq!( + message["tasks"].as_array().map(|tasks| tasks.len()), + Some(0), + "revoked ping capabilities must be applied before reconciliation" + ); + } + seen.insert(message_type.to_string()); + } + } + + assert_eq!( + seen, + expected.into_iter().map(str::to_string).collect(), + "capability changes should reconcile every agent desired-state domain" + ); let _ = sink.close().await; } diff --git a/crates/server/tests/integration.rs b/crates/server/tests/integration.rs index e41d9162a..d602f8f80 100644 --- a/crates/server/tests/integration.rs +++ b/crates/server/tests/integration.rs @@ -4451,9 +4451,9 @@ mod firewall_tests { serverbee_common::constants::PROTOCOL_VERSION, ) .await; - // First-connect path issues an extra BlocklistReset+BlocklistSync after - // the ack. Drain through both extras to leave the stream clean. - let _ = drain_through_ack(&mut reader, "pb-1", 2).await; + // AgentReady reconciliation follows the ack: ping, network, IP quality, + // then BlocklistReset+BlocklistSync. Drain all five messages. + let _ = drain_through_ack(&mut reader, "pb-1", 5).await; // Now POST a block and assert the agent receives BlocklistAdd. let resp = admin @@ -4572,7 +4572,7 @@ mod firewall_tests { ) .await; - let extras = drain_through_ack(&mut reader, "connect-sync", 2).await; + let extras = drain_through_ack(&mut reader, "connect-sync", 5).await; // Reset, then Sync are appended after the ack handler runs. let types: Vec<&str> = extras .iter() @@ -4609,7 +4609,7 @@ mod firewall_tests { serverbee_common::constants::PROTOCOL_VERSION, ) .await; - let _drain = drain_through_ack(&mut reader, "ack-1", 2).await; + let _drain = drain_through_ack(&mut reader, "ack-1", 5).await; let resp = admin .post(format!("{}/api/firewall/blocks", base_url)) @@ -4683,7 +4683,7 @@ mod firewall_tests { 1, ) .await; - let extras = drain_through_ack(&mut reader, "old-1", 0).await; + let extras = drain_through_ack(&mut reader, "old-1", 3).await; let types: Vec<&str> = extras .iter() .map(|m| m["type"].as_str().unwrap_or("")) diff --git a/crates/server/tests/router_probes_extra.rs b/crates/server/tests/router_probes_extra.rs index 7054e4ee4..d6ce85c0c 100644 --- a/crates/server/tests/router_probes_extra.rs +++ b/crates/server/tests/router_probes_extra.rs @@ -106,7 +106,6 @@ async fn network_probe_setting_update_pushes_sync_to_online_agent() { // Deleting a custom target that is assigned to an online server notifies that // agent with a fresh `NetworkProbeSync` (now without the deleted target). -// Exercises the delete_target → affected_server_ids → push branch. #[tokio::test] async fn network_probe_delete_assigned_target_pushes_sync() { let (base_url, _tmp) = start_test_server().await; @@ -168,6 +167,72 @@ async fn network_probe_delete_assigned_target_pushes_sync() { ); } +// Updating an assigned custom target must replace the agent's live projection. +// This is the regression that the old router-local fan-out missed entirely. +#[tokio::test] +async fn network_probe_update_assigned_target_pushes_latest_sync() { + let (base_url, _tmp) = start_test_server().await; + let admin = http_client(); + login_admin(&admin, &base_url).await; + + let (server_id, token) = register_agent(&admin, &base_url).await; + let (_sink, mut reader) = online_agent(&base_url, &token, None).await; + + let created: Value = admin + .post(format!("{}/api/network-probes/targets", base_url)) + .json(&json!({ + "name": "before-update", "provider": "P", "location": "L", + "target": "9.9.9.9", "probe_type": "icmp" + })) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + let target_id = created["data"]["id"].as_str().unwrap().to_string(); + + let assigned = admin + .put(format!( + "{}/api/servers/{}/network-probes/targets", + base_url, server_id + )) + .json(&json!({ "target_ids": [target_id] })) + .send() + .await + .unwrap(); + assert_eq!(assigned.status(), 200, "assigning targets succeeds"); + drain_initial_noise(&mut reader).await; + + let updated = admin + .put(format!( + "{}/api/network-probes/targets/{}", + base_url, target_id + )) + .json(&json!({ + "name": "after-update", + "target": "8.8.8.8", + "probe_type": "tcp" + })) + .send() + .await + .unwrap(); + assert_eq!(updated.status(), 200, "updating a custom target succeeds"); + + let sync = recv_until(&mut reader, "network_probe_sync").await; + let target = sync["targets"] + .as_array() + .and_then(|targets| { + targets + .iter() + .find(|target| target["target_id"].as_str() == Some(target_id.as_str())) + }) + .expect("updated target should remain assigned"); + assert_eq!(target["name"], "after-update"); + assert_eq!(target["target"], "8.8.8.8"); + assert_eq!(target["probe_type"], "tcp"); +} + // A preset target cannot be modified: update_target hits the Forbidden arm → 403. #[tokio::test] async fn network_probe_update_preset_target_forbidden_403() { From 775aa5c23626964582a95889f51de607b8e7c9f6 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 14:01:50 +0800 Subject: [PATCH 08/44] refactor(agent): own connection command state in one runtime handle_server_message threaded 13 borrows through every call and the test harness had to mirror that exact parameter list, so adding any manager meant touching the dispatcher signature, the select! loop, and the harness in lockstep. ConnectionRuntime now owns every manager, outbound sender, and piece of Docker lifecycle state a server command can touch; the dispatcher is a method on the state it drives. - reporter/runtime.rs owns the managers, capability authority, upgrade config, firewall manager, and Docker probe/retry/stats/demote cycle - the two Docker select! arms merge into one docker_tick() future, so the loop no longer juggles three loose Option/bool/Interval locals - the triple-duplicated connection-exit cleanup collapses into runtime.shutdown() - the dispatch test Harness builds a real runtime and calls the two-argument method instead of mirroring a 13-parameter signature wire and file_ops keep their existing seams; the dispatcher still delegates file replies to file_ops unchanged. --- crates/agent/src/reporter/mod.rs | 1694 +------------------------ crates/agent/src/reporter/runtime.rs | 1701 ++++++++++++++++++++++++++ 2 files changed, 1756 insertions(+), 1639 deletions(-) create mode 100644 crates/agent/src/reporter/runtime.rs diff --git a/crates/agent/src/reporter/mod.rs b/crates/agent/src/reporter/mod.rs index 07e31b6fd..c6676eb07 100644 --- a/crates/agent/src/reporter/mod.rs +++ b/crates/agent/src/reporter/mod.rs @@ -1,43 +1,35 @@ mod file_ops; +mod runtime; mod wire; use std::net::IpAddr; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; use futures_util::{SinkExt, StreamExt}; use rand::Rng; -use serverbee_common::constants::{DEFAULT_COMMAND_TIMEOUT_SECS, MAX_TASK_OUTPUT_SIZE}; use serverbee_common::protocol::{AgentMessage, ServerMessage, UpgradeStage}; -use serverbee_common::types::{NetworkInterface, NetworkProbeResultData}; +use serverbee_common::types::NetworkInterface; use sysinfo::Networks; use tokio::sync::mpsc; use tokio::time::interval; use tokio_tungstenite::connect_async; use tokio_tungstenite::tungstenite::Message; +use runtime::{ConnectionRuntime, DockerTick}; use wire::send_msg; use crate::capability_grants::CapabilityAuthority; use crate::collector::Collector; use crate::config::AgentConfig; -use crate::docker::DockerManager; -use crate::file_manager::{FileEvent, FileManager}; use crate::firewall::FirewallManager; use crate::firewall::nft::CliNftExecutor; -use crate::ip_quality::{RunResult, UnlockChecker}; -use crate::network_prober::NetworkProber; -use crate::pinger::PingManager; use crate::register; -use crate::terminal::{TerminalEvent, TerminalManager}; +use crate::terminal::TerminalEvent; const MAX_BACKOFF_SECS: u64 = 30; const JITTER_FACTOR: f64 = 0.2; const MAX_REREGISTER_ATTEMPTS: u32 = 3; -const DOCKER_RETRY_SECS: u64 = 30; - -static UPGRADE_IN_PROGRESS: AtomicBool = AtomicBool::new(false); pub struct Reporter { config: AgentConfig, @@ -188,41 +180,16 @@ impl Reporter { None => anyhow::bail!("Connection closed before Welcome"), }; - // Docker manager setup - let (docker_tx, mut docker_rx) = mpsc::channel::(256); - let mut docker_manager: Option = None; - let mut docker_available = false; - - // Try to initialize Docker connection - match DockerManager::try_new(docker_tx.clone(), Arc::clone(&capabilities)) { - Ok(dm) => match dm.verify_connection().await { - Ok(()) => { - tracing::info!("Docker daemon connected"); - docker_available = true; - docker_manager = Some(dm); - } - Err(e) => { - tracing::info!("Docker daemon not available: {e}"); - } - }, - Err(e) => { - tracing::info!("Docker not available: {e}"); - } - } - - // Docker retry interval (only used when docker_manager is None) - let mut docker_retry_interval = interval(Duration::from_secs(DOCKER_RETRY_SECS)); - docker_retry_interval.tick().await; // consume immediate tick - - // Separate stats interval (managed by start/stop stats commands) - // Uses a long default that gets replaced when stats are requested. - let mut docker_stats_interval: Option = None; - - // Build features list - let mut features = Vec::new(); - if docker_available { - features.push("docker".to_string()); - } + // Connection-scoped command runtime: owns every manager and channel + // the server-message dispatcher drives. Docker detection runs before + // SystemInfo so the advertised features list is accurate. + let (mut runtime, mut events) = ConnectionRuntime::new( + &self.config, + Arc::clone(&capabilities), + Arc::clone(&self.firewall_manager), + ); + runtime.probe_docker().await; + let features = runtime.features(); // Send SystemInfo let mut collector = Collector::new( @@ -257,22 +224,6 @@ impl Reporter { send_msg(&mut write, &info_msg).await?; tracing::info!("Sent SystemInfo"); - // Ping probe manager - let (ping_tx, mut ping_rx) = mpsc::channel(256); - let mut ping_manager = PingManager::new(ping_tx, Arc::clone(&capabilities)); - - // Terminal session manager - let (term_tx, mut term_rx) = mpsc::channel(256); - let mut terminal_manager = TerminalManager::new(term_tx, Arc::clone(&capabilities)); - - // Network probe manager - let (network_probe_tx, mut network_probe_rx) = mpsc::channel::(256); - let mut network_prober = NetworkProber::new(network_probe_tx, Arc::clone(&capabilities)); - - // IP quality unlock checker - let (unlock_result_tx, mut unlock_result_rx) = mpsc::channel::(8); - let unlock_checker = UnlockChecker::new(Arc::clone(&capabilities), unlock_result_tx); - // Seed the checker with the interface-derived public IP so the very first // run has a non-empty egress_ip even on stable VPS deployments where the // external IP never "changes" (i.e. IpChanged is never emitted because @@ -287,17 +238,14 @@ impl Reporter { // If neither is available the checker stays at None and will be // populated by IpChanged from spawn_external_ip_refresh. if seed_ip.is_some() { - unlock_checker.notify_ip_changed(seed_ip); + runtime.unlock_checker.notify_ip_changed(seed_ip); } } } - // File manager - let (file_tx, mut file_rx) = mpsc::channel::(16); - let file_manager = FileManager::new(self.config.file.clone(), Arc::clone(&capabilities)); - - // Channel for background command execution results. - let (cmd_result_tx, mut cmd_result_rx) = mpsc::channel::(32); + // Outbound channel for background command results; also feeds the + // fire-and-forget IP refresh tasks below. + let cmd_result_tx = runtime.cmd_result_tx.clone(); // Capability transitions from the process-wide authority: each one is // forwarded to the server as CapabilitiesChanged, after reconciling @@ -347,19 +295,19 @@ impl Reporter { send_msg(&mut write, &msg).await?; tracing::debug!("Sent report"); } - Some(ping_result) = ping_rx.recv() => { + Some(ping_result) = events.ping_rx.recv() => { let msg = AgentMessage::PingResult(ping_result); send_msg(&mut write, &msg).await?; tracing::debug!("Sent PingResult"); } - Some(cmd_msg) = cmd_result_rx.recv() => { + Some(cmd_msg) = events.cmd_result_rx.recv() => { // Intercept IpChanged to notify the UnlockChecker before forwarding. if let AgentMessage::IpChanged { ref ipv4, ref ipv6, .. } = cmd_msg { use serverbee_common::constants::CAP_IP_QUALITY; if capabilities.has(CAP_IP_QUALITY) { // Prefer IPv4 egress; fall back to IPv6. let new_ip = ipv4.clone().or_else(|| ipv6.clone()); - unlock_checker.notify_ip_changed(new_ip); + runtime.unlock_checker.notify_ip_changed(new_ip); } } send_msg(&mut write, &cmd_msg).await?; @@ -372,7 +320,7 @@ impl Reporter { // revoked or expired terminal capability must tear // down live PTY sessions, not just block new ones. if !has_capability(t.effective, CAP_TERMINAL) { - for session_id in terminal_manager.close_all() { + for session_id in runtime.terminal_manager.close_all() { let msg = AgentMessage::TerminalError { session_id, error: "Terminal capability revoked or expired".to_string(), @@ -409,7 +357,7 @@ impl Reporter { } } } - Some(run_result) = unlock_result_rx.recv() => { + Some(run_result) = events.unlock_result_rx.recv() => { let msg = AgentMessage::UnlockResults { egress_ip: run_result.egress_ip, results: run_result.results, @@ -427,7 +375,7 @@ impl Reporter { send_msg(&mut write, &external_msg).await?; tracing::debug!("Sent external agent message"); } - Some(term_event) = term_rx.recv() => { + Some(term_event) = events.term_rx.recv() => { let msg = match term_event { TerminalEvent::Output { session_id, data } => { AgentMessage::TerminalOutput { session_id, data } @@ -439,7 +387,7 @@ impl Reporter { AgentMessage::TerminalError { session_id, error } } TerminalEvent::Exited { session_id } => { - terminal_manager.close(&session_id); + runtime.terminal_manager.close(&session_id); AgentMessage::TerminalError { session_id, error: "Session exited".to_string(), @@ -448,10 +396,10 @@ impl Reporter { }; send_msg(&mut write, &msg).await?; } - Some(first_result) = network_probe_rx.recv() => { + Some(first_result) = events.network_probe_rx.recv() => { let mut results = vec![first_result]; // Drain any additional results that arrived at the same time - while let Ok(additional) = network_probe_rx.try_recv() { + while let Ok(additional) = events.network_probe_rx.try_recv() { results.push(additional); } let count = results.len(); @@ -459,34 +407,22 @@ impl Reporter { send_msg(&mut write, &msg).await?; tracing::debug!("Sent NetworkProbeResults ({count} results)"); } - Some(file_event) = file_rx.recv() => { + Some(file_event) = events.file_rx.recv() => { let msg: AgentMessage = file_event.into(); send_msg(&mut write, &msg).await?; tracing::debug!("Sent file event"); } // Docker messages from DockerManager background tasks - Some(docker_msg) = docker_rx.recv() => { + Some(docker_msg) = events.docker_rx.recv() => { send_msg(&mut write, &docker_msg).await?; tracing::debug!("Sent Docker message"); } - // Docker stats polling (uses separate interval to avoid borrow conflicts) - Some(_) = async { - match docker_stats_interval.as_mut() { - Some(iv) => Some(iv.tick().await), - None => None, - } - } => { - if let Some(dm) = docker_manager.as_mut() - && let Err(e) = dm.poll_stats().await - { - tracing::warn!("Docker stats polling failed: {e}"); - self.demote_docker_runtime( - &mut write, - &mut docker_manager, - &mut docker_available, - &mut docker_stats_interval, - ) - .await?; + // Docker housekeeping: stats polling while the daemon is + // connected, reconnect retry while it is absent. + tick = runtime.docker_tick() => { + match tick { + DockerTick::PollStats => runtime.poll_docker_stats(&mut write).await?, + DockerTick::Retry => runtime.retry_docker(&mut write).await?, } } // IP change detection — spawned off the WS hot path so a @@ -512,47 +448,14 @@ impl Reporter { ); } } - // Docker retry (reconnect when docker is unavailable) - _ = docker_retry_interval.tick(), if docker_manager.is_none() => { - tracing::debug!("Retrying Docker connection..."); - match DockerManager::try_new(docker_tx.clone(), Arc::clone(&capabilities)) { - Ok(dm) => { - match dm.verify_connection().await { - Ok(()) => { - tracing::info!("Docker daemon now available"); - docker_manager = Some(dm); - docker_available = true; - // Notify server about features change - let msg = AgentMessage::FeaturesUpdate { - features: vec!["docker".to_string()], - }; - send_msg(&mut write, &msg).await?; - } - Err(e) => { - tracing::debug!("Docker still not available: {e}"); - } - } - } - Err(e) => { - tracing::debug!("Docker still not available: {e}"); - } - } - } server_msg = read.next() => { match server_msg { Some(Ok(Message::Text(text))) => { - self.handle_server_message(&text, &mut write, &mut ping_manager, &mut terminal_manager, &mut network_prober, &cmd_result_tx, &capabilities, &file_manager, &file_tx, &mut docker_manager, &mut docker_available, &mut docker_stats_interval, &unlock_checker).await?; + runtime.handle_server_message(&text, &mut write).await?; } Some(Ok(Message::Close(_))) => { tracing::info!("Server closed connection"); - ping_manager.stop_all(); - terminal_manager.close_all(); - network_prober.stop_all(); - unlock_checker.stop(); - file_manager.cancel_all_transfers(); - if let Some(dm) = docker_manager.as_mut() { - dm.cleanup(); - } + runtime.shutdown(); return Ok(()); } Some(Ok(Message::Ping(data))) => { @@ -561,26 +464,12 @@ impl Reporter { Some(Ok(_)) => {} Some(Err(e)) => { tracing::error!("WebSocket error: {e}"); - ping_manager.stop_all(); - terminal_manager.close_all(); - network_prober.stop_all(); - unlock_checker.stop(); - file_manager.cancel_all_transfers(); - if let Some(dm) = docker_manager.as_mut() { - dm.cleanup(); - } + runtime.shutdown(); return Err(e.into()); } None => { tracing::info!("WebSocket stream ended"); - ping_manager.stop_all(); - terminal_manager.close_all(); - network_prober.stop_all(); - unlock_checker.stop(); - file_manager.cancel_all_transfers(); - if let Some(dm) = docker_manager.as_mut() { - dm.cleanup(); - } + runtime.shutdown(); return Ok(()); } } @@ -588,434 +477,6 @@ impl Reporter { } } } - - #[allow(clippy::too_many_arguments)] - async fn handle_server_message( - &mut self, - text: &str, - write: &mut S, - ping_manager: &mut PingManager, - terminal_manager: &mut TerminalManager, - network_prober: &mut NetworkProber, - cmd_result_tx: &mpsc::Sender, - capabilities: &CapabilityAuthority, - file_manager: &FileManager, - file_tx: &mpsc::Sender, - docker_manager: &mut Option, - docker_available: &mut bool, - docker_stats_interval: &mut Option, - unlock_checker: &UnlockChecker, - ) -> anyhow::Result<()> - where - S: SinkExt + Unpin, - { - use serverbee_common::constants::*; - - let msg: ServerMessage = match serde_json::from_str(text) { - Ok(m) => m, - Err(e) => { - tracing::warn!("Failed to parse server message: {e}"); - return Ok(()); - } - }; - - match msg { - ServerMessage::Ping => { - send_msg(write, &AgentMessage::Pong).await?; - tracing::debug!("Responded to Ping with Pong"); - } - ServerMessage::Exec { - task_id, - command, - timeout, - } => { - if !capabilities.has(CAP_EXEC) { - tracing::warn!("Exec denied: capability disabled (task_id={task_id})"); - spawn_capability_denied(cmd_result_tx, Some(task_id), "exec"); - return Ok(()); - } - tracing::info!("Executing command (task_id={task_id}): {command}"); - let tx = cmd_result_tx.clone(); - tokio::spawn(async move { - let result = execute_command(&task_id, &command, timeout).await; - let msg = AgentMessage::TaskResult { - msg_id: uuid::Uuid::new_v4().to_string(), - result, - }; - if tx.send(msg).await.is_err() { - tracing::warn!( - "Failed to send TaskResult for task_id={task_id}: channel closed" - ); - } else { - tracing::info!("TaskResult ready for task_id={task_id}"); - } - }); - } - ServerMessage::Ack { msg_id } => { - tracing::debug!("Received Ack for msg_id={msg_id}"); - } - ServerMessage::Welcome { .. } => { - tracing::warn!("Unexpected second Welcome message"); - } - ServerMessage::PingTasksSync { tasks } => { - tracing::info!("Received PingTasksSync with {} tasks", tasks.len()); - ping_manager.sync(tasks); - } - ServerMessage::TerminalOpen { - session_id, - rows, - cols, - } => { - tracing::info!("Opening terminal session {session_id} ({cols}x{rows})"); - terminal_manager.open(session_id, rows, cols); - } - ServerMessage::TerminalInput { session_id, data } => { - terminal_manager.write_input(&session_id, &data); - } - ServerMessage::TerminalResize { - session_id, - rows, - cols, - } => { - tracing::debug!("Resizing terminal {session_id} to {cols}x{rows}"); - terminal_manager.resize(&session_id, rows, cols); - } - ServerMessage::TerminalClose { session_id } => { - tracing::debug!("Closing terminal session {session_id}"); - terminal_manager.close(&session_id); - } - ServerMessage::Upgrade { - version, - job_id, - .. - } => { - if !capabilities.has(CAP_UPGRADE) { - tracing::warn!("Upgrade denied: capability disabled"); - send_msg(write, &capability_denied_msg(None, "upgrade")).await?; - return Ok(()); - } - - if UPGRADE_IN_PROGRESS - .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) - .is_err() - { - let tx = cmd_result_tx.clone(); - tokio::spawn(async move { - emit_upgrade_failure( - &tx, - job_id, - version, - UpgradeStage::Downloading, - "another upgrade is already running".to_string(), - None, - ) - .await; - }); - return Ok(()); - } - - tracing::info!("Upgrade requested: v{version} (pinned source)"); - let upgrade_cfg = self.config.upgrade.clone(); - let tx = cmd_result_tx.clone(); - tokio::spawn(async move { - if let Err(e) = - perform_upgrade(&version, &upgrade_cfg, job_id, tx.clone()).await - { - tracing::error!("Upgrade to v{version} failed: {e}"); - UPGRADE_IN_PROGRESS.store(false, Ordering::SeqCst); - } - }); - } - ServerMessage::NetworkProbeSync { - targets, - interval, - packet_count, - } => { - tracing::info!( - "Received NetworkProbeSync: {} targets, interval={}s, packet_count={}", - targets.len(), - interval, - packet_count - ); - network_prober.sync(targets, interval, packet_count); - } - ServerMessage::Traceroute { - request_id, - target, - max_hops, - protocol, - } => { - if !capabilities.has(CAP_PING_ICMP) { - tracing::warn!( - "Traceroute denied: capability disabled (request_id={request_id})" - ); - spawn_capability_denied(cmd_result_tx, Some(request_id), "ping_icmp"); - return Ok(()); - } - - // Input validation: target must be domain or IP only. - if !crate::traceroute::is_valid_traceroute_target(&target) { - tracing::warn!( - "Traceroute rejected: invalid target '{target}' (request_id={request_id})" - ); - let tx = cmd_result_tx.clone(); - let request_id_c = request_id.clone(); - let target_c = target.clone(); - tokio::spawn(async move { - let _ = tx - .send(AgentMessage::TracerouteRoundUpdate { - request_id: request_id_c, - target: target_c, - round: 0, - total_rounds: 0, - hops: vec![], - completed: true, - error: Some( - "Invalid target: must be a domain or IP address".into(), - ), - }) - .await; - }); - return Ok(()); - } - - let proto = - protocol.unwrap_or(serverbee_common::protocol::TraceProtocol::Icmp); - tracing::info!( - "Executing traceroute to {target} (max_hops={max_hops}, request_id={request_id}, protocol={proto:?})" - ); - crate::traceroute::spawn_traceroute( - request_id, - target, - max_hops, - proto, - cmd_result_tx.clone(), - ); - } - // --- File management messages --- (gate + reply shapes live in file_ops) - msg @ (ServerMessage::FileList { .. } - | ServerMessage::FileStat { .. } - | ServerMessage::FileRead { .. } - | ServerMessage::FileWrite { .. } - | ServerMessage::FileDelete { .. } - | ServerMessage::FileMkdir { .. } - | ServerMessage::FileMove { .. } - | ServerMessage::FileDownloadStart { .. } - | ServerMessage::FileDownloadCancel { .. } - | ServerMessage::FileUploadStart { .. } - | ServerMessage::FileUploadChunk { .. } - | ServerMessage::FileUploadEnd { .. }) => { - file_ops::handle_file_message(msg, write, file_manager, file_tx, capabilities) - .await?; - } - // --- Docker messages --- - ServerMessage::DockerStartStats { interval_secs } => { - if docker_manager.is_some() { - let secs = interval_secs.max(1); - tracing::info!("Starting Docker stats polling every {secs}s"); - let mut iv = tokio::time::interval(Duration::from_secs(secs as u64)); - iv.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - *docker_stats_interval = Some(iv); - } else { - tracing::warn!("DockerStartStats received but Docker is not available"); - let unavailable = AgentMessage::DockerUnavailable { msg_id: None }; - send_msg(write, &unavailable).await?; - } - } - ServerMessage::DockerStopStats => { - tracing::info!("Stopping Docker stats polling"); - *docker_stats_interval = None; - } - ServerMessage::DockerListContainers { .. } - | ServerMessage::DockerLogsStart { .. } - | ServerMessage::DockerLogsStop { .. } - | ServerMessage::DockerEventsStart - | ServerMessage::DockerEventsStop - | ServerMessage::DockerContainerAction { .. } - | ServerMessage::DockerGetInfo { .. } - | ServerMessage::DockerListNetworks { .. } - | ServerMessage::DockerListVolumes { .. } => { - if let Some(dm) = docker_manager.as_mut() { - if let Err(e) = dm.handle_server_message(msg.clone()).await { - tracing::warn!("Docker runtime became unavailable: {e}"); - self.demote_docker_runtime( - write, - docker_manager, - docker_available, - docker_stats_interval, - ) - .await?; - } - } else { - tracing::warn!("Docker message received but Docker is not available"); - let unavailable = AgentMessage::DockerUnavailable { - msg_id: docker_request_msg_id(&msg), - }; - send_msg(write, &unavailable).await?; - } - } - // Firewall blocklist variants — dispatched to the FirewallManager - // state machine; any returned ack is sent straight back over the - // WebSocket. - // - // The mutating variants (Sync/Add/Remove) enforce CAP_FIREWALL_BLOCK - // on the agent's own host, mirroring the capability gates on Exec / - // File / Traceroute etc. — the server is not the only trust boundary. - // BlocklistReset is deliberately *not* gated: it wipes ServerBee's - // own nft table (cleanup / disable path) and must stay reachable even - // after the capability is revoked, so a denied agent can still be - // cleaned up. - msg @ (ServerMessage::BlocklistSync { .. } - | ServerMessage::BlocklistAdd { .. } - | ServerMessage::BlocklistRemove { .. } - | ServerMessage::BlocklistReset) => { - let is_reset = matches!(msg, ServerMessage::BlocklistReset); - if !is_reset && !capabilities.has(CAP_FIREWALL_BLOCK) { - tracing::warn!( - "Firewall blocklist mutation denied: CAP_FIREWALL_BLOCK not effective — ignoring" - ); - } else if let Some(reply) = self.firewall_manager.handle(msg).await { - send_msg(write, &reply).await?; - tracing::debug!("Sent firewall blocklist ack"); - } - } - ServerMessage::IpQualitySync { services, interval_hours } => { - if capabilities.has(CAP_IP_QUALITY) { - tracing::info!( - "Received IpQualitySync: {} services, interval={}h", - services.len(), - interval_hours - ); - unlock_checker.sync(services, interval_hours).await; - } else { - tracing::debug!( - "IpQualitySync received but CAP_IP_QUALITY not effective — ignoring" - ); - } - } - ServerMessage::IpQualityRunNow => { - if capabilities.has(CAP_IP_QUALITY) { - tracing::info!("Received IpQualityRunNow"); - unlock_checker.run_now(); - } else { - tracing::debug!( - "IpQualityRunNow received but CAP_IP_QUALITY not effective — ignoring" - ); - } - } - } - - Ok(()) - } -} - -impl Reporter { - async fn demote_docker_runtime( - &self, - write: &mut S, - docker_manager: &mut Option, - docker_available: &mut bool, - docker_stats_interval: &mut Option, - ) -> anyhow::Result<()> - where - S: SinkExt + Unpin, - { - if let Some(dm) = docker_manager.as_mut() { - dm.cleanup(); - } - *docker_manager = None; - *docker_stats_interval = None; - - if *docker_available { - *docker_available = false; - let msg = AgentMessage::FeaturesUpdate { features: vec![] }; - send_msg(write, &msg).await?; - } - - Ok(()) - } -} - -/// Build the standard agent-local capability denial message. -fn capability_denied_msg(msg_id: Option, capability: &str) -> AgentMessage { - AgentMessage::CapabilityDenied { - msg_id, - session_id: None, - capability: capability.to_string(), - reason: serverbee_common::constants::CapabilityDeniedReason::AgentCapabilityDisabled, - } -} - -/// Queue a capability denial onto the outbound channel without blocking the -/// read loop. -fn spawn_capability_denied( - tx: &mpsc::Sender, - msg_id: Option, - capability: &str, -) { - let denied = capability_denied_msg(msg_id, capability); - let tx = tx.clone(); - tokio::spawn(async move { - let _ = tx.send(denied).await; - }); -} - -fn docker_request_msg_id(msg: &ServerMessage) -> Option { - match msg { - ServerMessage::DockerListContainers { msg_id } - | ServerMessage::DockerGetInfo { msg_id } - | ServerMessage::DockerListNetworks { msg_id } - | ServerMessage::DockerListVolumes { msg_id } => Some(msg_id.clone()), - ServerMessage::DockerContainerAction { msg_id, .. } => Some(msg_id.clone()), - _ => None, - } -} - -async fn execute_command( - task_id: &str, - command: &str, - timeout: Option, -) -> serverbee_common::types::TaskResult { - let timeout_secs = timeout.unwrap_or(DEFAULT_COMMAND_TIMEOUT_SECS); - - let result = tokio::time::timeout( - Duration::from_secs(timeout_secs as u64), - tokio::process::Command::new("sh") - .arg("-c") - .arg(command) - .output(), - ) - .await; - - match result { - Ok(Ok(output)) => { - let mut combined = String::from_utf8_lossy(&output.stdout).to_string(); - let stderr = String::from_utf8_lossy(&output.stderr); - if !stderr.is_empty() { - combined.push('\n'); - combined.push_str(&stderr); - } - if combined.len() > MAX_TASK_OUTPUT_SIZE { - combined.truncate(MAX_TASK_OUTPUT_SIZE); - combined.push_str("\n... (output truncated)"); - } - serverbee_common::types::TaskResult { - task_id: task_id.to_string(), - output: combined, - exit_code: output.status.code().unwrap_or(-1), - } - } - Ok(Err(e)) => serverbee_common::types::TaskResult { - task_id: task_id.to_string(), - output: format!("Failed to execute command: {e}"), - exit_code: -1, - }, - Err(_) => serverbee_common::types::TaskResult { - task_id: task_id.to_string(), - output: format!("Command timed out after {timeout_secs}s"), - exit_code: -1, - }, - } } fn build_ws_url(config: &AgentConfig) -> anyhow::Result { @@ -1469,10 +930,7 @@ mod tests { assert_eq!(mk("https://example.com"), "wss://example.com/api/agent/ws"); assert_eq!(mk("http://example.com"), "ws://example.com/api/agent/ws"); // trailing slash trimmed before appending the path - assert_eq!( - mk("https://example.com/"), - "wss://example.com/api/agent/ws" - ); + assert_eq!(mk("https://example.com/"), "wss://example.com/api/agent/ws"); // schemeless host gets ws:// prefix assert_eq!(mk("example.com:9527"), "ws://example.com:9527/api/agent/ws"); } @@ -1547,1060 +1005,6 @@ mod tests { ); } - #[test] - fn test_docker_request_msg_id_extracts_only_request_variants() { - // Variants that carry a msg_id return it... - assert_eq!( - docker_request_msg_id(&ServerMessage::DockerListContainers { - msg_id: "a".to_string() - }), - Some("a".to_string()) - ); - assert_eq!( - docker_request_msg_id(&ServerMessage::DockerGetInfo { - msg_id: "b".to_string() - }), - Some("b".to_string()) - ); - assert_eq!( - docker_request_msg_id(&ServerMessage::DockerListNetworks { - msg_id: "c".to_string() - }), - Some("c".to_string()) - ); - assert_eq!( - docker_request_msg_id(&ServerMessage::DockerListVolumes { - msg_id: "d".to_string() - }), - Some("d".to_string()) - ); - assert_eq!( - docker_request_msg_id(&ServerMessage::DockerContainerAction { - msg_id: "e".to_string(), - container_id: "cid".to_string(), - action: serverbee_common::docker_types::DockerAction::Restart { timeout: None }, - }), - Some("e".to_string()) - ); - // ...non-request docker variants return None. - assert_eq!( - docker_request_msg_id(&ServerMessage::DockerStopStats), - None - ); - assert_eq!(docker_request_msg_id(&ServerMessage::Ping), None); - } - - // ---------------------------------------------------------------------- - // `handle_server_message` dispatcher coverage via a mock sink. - // - // The dispatcher is generic over the WS sink (`S: SinkExt`), - // so we drive it with an in-memory recording sink instead of a real - // WebSocket. A `Harness` owns every manager + channel the method borrows, - // keeping receivers alive so spawned senders never error. - // ---------------------------------------------------------------------- - - /// In-memory sink that records every `Message` written to it. All poll_* - /// hooks succeed immediately; `start_send` just pushes into a shared Vec. - #[derive(Clone)] - struct RecordingSink { - sent: Arc>>, - } - - impl RecordingSink { - fn new() -> Self { - Self { - sent: Arc::new(std::sync::Mutex::new(Vec::new())), - } - } - - /// All recorded messages decoded into `AgentMessage` (text frames only). - fn agent_messages(&self) -> Vec { - self.sent - .lock() - .unwrap() - .iter() - .filter_map(|m| match m { - Message::Text(t) => serde_json::from_str::(t.as_str()).ok(), - _ => None, - }) - .collect() - } - - fn sent_count(&self) -> usize { - self.sent.lock().unwrap().len() - } - } - - impl futures_util::Sink for RecordingSink { - type Error = tokio_tungstenite::tungstenite::Error; - - fn poll_ready( - self: std::pin::Pin<&mut Self>, - _cx: &mut std::task::Context<'_>, - ) -> std::task::Poll> { - std::task::Poll::Ready(Ok(())) - } - - fn start_send( - self: std::pin::Pin<&mut Self>, - item: Message, - ) -> Result<(), Self::Error> { - self.sent.lock().unwrap().push(item); - Ok(()) - } - - fn poll_flush( - self: std::pin::Pin<&mut Self>, - _cx: &mut std::task::Context<'_>, - ) -> std::task::Poll> { - std::task::Poll::Ready(Ok(())) - } - - fn poll_close( - self: std::pin::Pin<&mut Self>, - _cx: &mut std::task::Context<'_>, - ) -> std::task::Poll> { - std::task::Poll::Ready(Ok(())) - } - } - - /// Owns every manager + channel borrowed by `handle_server_message`, - /// plus the receiver ends so background senders never see a closed channel. - struct Harness { - reporter: Reporter, - capabilities: Arc, - ping_manager: PingManager, - terminal_manager: TerminalManager, - network_prober: NetworkProber, - cmd_result_tx: mpsc::Sender, - cmd_result_rx: mpsc::Receiver, - file_manager: FileManager, - file_tx: mpsc::Sender, - docker_manager: Option, - docker_available: bool, - docker_stats_interval: Option, - unlock_checker: UnlockChecker, - // Keep manager-side receivers alive for the test's lifetime. - _ping_rx: mpsc::Receiver, - _term_rx: mpsc::Receiver, - _network_rx: mpsc::Receiver, - _file_rx: mpsc::Receiver, - _unlock_rx: mpsc::Receiver, - } - - impl Harness { - /// Build a harness with all capability bits ON by default. `file_cfg` - /// lets individual tests enable the file manager with a temp root. - fn new(caps: u32, file_cfg: FileConfig) -> Self { - let config = AgentConfig { - server_url: "http://127.0.0.1:9527".to_string(), - token: "t".to_string(), - enrollment_code: String::new(), - collector: CollectorConfig::default(), - log: LogConfig::default(), - file: file_cfg.clone(), - ip_change: IpChangeConfig::default(), - upgrade: UpgradeConfig::default(), - security: SecurityConfig::default(), - capabilities: CapabilitiesConfig::default(), - }; - let capabilities = CapabilityAuthority::fixed(caps); - let reporter = Reporter::new(config, "fp".to_string(), Arc::clone(&capabilities)); - - let (ping_tx, _ping_rx) = mpsc::channel(16); - let ping_manager = PingManager::new(ping_tx, Arc::clone(&capabilities)); - - let (term_tx, _term_rx) = mpsc::channel(16); - let terminal_manager = TerminalManager::new(term_tx, Arc::clone(&capabilities)); - - let (network_tx, _network_rx) = mpsc::channel(16); - let network_prober = NetworkProber::new(network_tx, Arc::clone(&capabilities)); - - let (cmd_result_tx, cmd_result_rx) = mpsc::channel(32); - - let (file_tx, _file_rx) = mpsc::channel(16); - let file_manager = FileManager::new(file_cfg, Arc::clone(&capabilities)); - - let (unlock_tx, _unlock_rx) = mpsc::channel(8); - let unlock_checker = UnlockChecker::new(Arc::clone(&capabilities), unlock_tx); - - Self { - reporter, - capabilities, - ping_manager, - terminal_manager, - network_prober, - cmd_result_tx, - cmd_result_rx, - file_manager, - file_tx, - docker_manager: None, - docker_available: false, - docker_stats_interval: None, - unlock_checker, - _ping_rx, - _term_rx, - _network_rx, - _file_rx, - _unlock_rx, - } - } - - /// Dispatch a single `ServerMessage` (as JSON) through the method. - async fn dispatch(&mut self, text: &str, sink: &mut RecordingSink) -> anyhow::Result<()> { - self.reporter - .handle_server_message( - text, - sink, - &mut self.ping_manager, - &mut self.terminal_manager, - &mut self.network_prober, - &self.cmd_result_tx, - &self.capabilities, - &self.file_manager, - &self.file_tx, - &mut self.docker_manager, - &mut self.docker_available, - &mut self.docker_stats_interval, - &self.unlock_checker, - ) - .await - } - } - - /// All capability bits set — every success arm runs. - const ALL_CAPS: u32 = serverbee_common::constants::CAP_VALID_MASK; - - fn enabled_file_cfg(root: &std::path::Path) -> FileConfig { - FileConfig { - enabled: true, - root_paths: vec![root.to_string_lossy().to_string()], - ..FileConfig::default() - } - } - - #[tokio::test] - async fn test_blocklist_mutation_denied_without_firewall_capability() { - use serverbee_common::constants::CAP_FIREWALL_BLOCK; - // All caps except firewall block — simulates a revoked capability. - let caps = ALL_CAPS & !CAP_FIREWALL_BLOCK; - let mut h = Harness::new(caps, FileConfig::default()); - let mut sink = RecordingSink::new(); - - // A mutating variant must be dropped before reaching the firewall - // manager, so nothing is written back and no nft command runs. - h.dispatch(r#"{"type":"blocklist_remove","id":"x"}"#, &mut sink) - .await - .unwrap(); - assert_eq!( - sink.sent_count(), - 0, - "blocklist mutation must be ignored when CAP_FIREWALL_BLOCK is off" - ); - } - - #[tokio::test] - async fn test_dispatch_unparseable_text_is_ignored() { - let mut h = Harness::new(ALL_CAPS, FileConfig::default()); - let mut sink = RecordingSink::new(); - // Not valid JSON for ServerMessage — must be swallowed as Ok with no output. - h.dispatch("this is not json", &mut sink).await.unwrap(); - h.dispatch(r#"{"type":"nonexistent_variant"}"#, &mut sink) - .await - .unwrap(); - assert_eq!(sink.sent_count(), 0, "unparseable input must not emit anything"); - } - - #[tokio::test] - async fn test_dispatch_ping_responds_with_pong() { - let mut h = Harness::new(ALL_CAPS, FileConfig::default()); - let mut sink = RecordingSink::new(); - h.dispatch(r#"{"type":"ping"}"#, &mut sink).await.unwrap(); - let msgs = sink.agent_messages(); - assert_eq!(msgs.len(), 1); - assert!(matches!(msgs[0], AgentMessage::Pong)); - } - - #[tokio::test] - async fn test_dispatch_ack_and_welcome_are_noops() { - let mut h = Harness::new(ALL_CAPS, FileConfig::default()); - let mut sink = RecordingSink::new(); - h.dispatch(r#"{"type":"ack","msg_id":"m1"}"#, &mut sink) - .await - .unwrap(); - h.dispatch( - r#"{"type":"welcome","server_id":"s","protocol_version":1,"report_interval":3}"#, - &mut sink, - ) - .await - .unwrap(); - assert_eq!(sink.sent_count(), 0, "ack/welcome must not write to the sink"); - } - - #[tokio::test] - async fn test_dispatch_exec_denied_when_capability_absent() { - // CAP_EXEC missing -> a CapabilityDenied is pushed onto cmd_result_tx - // (NOT the sink). We drain the channel to assert. - let caps = ALL_CAPS & !serverbee_common::constants::CAP_EXEC; - let mut h = Harness::new(caps, FileConfig::default()); - let mut sink = RecordingSink::new(); - h.dispatch( - r#"{"type":"exec","task_id":"task-42","command":"true","timeout":1}"#, - &mut sink, - ) - .await - .unwrap(); - assert_eq!(sink.sent_count(), 0, "denied exec writes to channel, not sink"); - let denied = h.cmd_result_rx.recv().await.expect("denied msg expected"); - match denied { - AgentMessage::CapabilityDenied { - msg_id, - capability, - reason, - .. - } => { - assert_eq!(msg_id, Some("task-42".to_string())); - assert_eq!(capability, "exec"); - assert_eq!(reason, CapabilityDeniedReason::AgentCapabilityDisabled); - } - other => panic!("expected CapabilityDenied, got {other:?}"), - } - } - - #[tokio::test] - async fn test_dispatch_exec_allowed_runs_and_emits_task_result() { - // `true` is a deterministic, always-present POSIX builtin/command. - let mut h = Harness::new(ALL_CAPS, FileConfig::default()); - let mut sink = RecordingSink::new(); - h.dispatch( - r#"{"type":"exec","task_id":"task-ok","command":"true","timeout":5}"#, - &mut sink, - ) - .await - .unwrap(); - // The execution is spawned; await its TaskResult on the channel. - let result = tokio::time::timeout(Duration::from_secs(10), h.cmd_result_rx.recv()) - .await - .expect("task did not complete in time") - .expect("TaskResult expected"); - match result { - AgentMessage::TaskResult { result, .. } => { - assert_eq!(result.task_id, "task-ok"); - assert_eq!(result.exit_code, 0); - } - other => panic!("expected TaskResult, got {other:?}"), - } - } - - #[tokio::test] - async fn test_dispatch_ping_tasks_sync_is_accepted() { - let mut h = Harness::new(ALL_CAPS, FileConfig::default()); - let mut sink = RecordingSink::new(); - // Empty task list — sync clears the manager; dispatcher returns Ok with - // no WS output. - h.dispatch(r#"{"type":"ping_tasks_sync","tasks":[]}"#, &mut sink) - .await - .unwrap(); - assert_eq!(sink.sent_count(), 0); - } - - #[tokio::test] - async fn test_dispatch_network_probe_sync_is_accepted() { - let mut h = Harness::new(ALL_CAPS, FileConfig::default()); - let mut sink = RecordingSink::new(); - h.dispatch( - r#"{"type":"network_probe_sync","targets":[],"interval":30,"packet_count":3}"#, - &mut sink, - ) - .await - .unwrap(); - assert_eq!(sink.sent_count(), 0); - } - - #[tokio::test] - async fn test_dispatch_terminal_lifecycle_without_capability() { - // CAP_TERMINAL off: open() routes to the denied event (no PTY spawned), - // and input/resize/close on a missing session are safe no-ops. The - // dispatcher must return Ok for each and never touch the sink. - let caps = ALL_CAPS & !serverbee_common::constants::CAP_TERMINAL; - let mut h = Harness::new(caps, FileConfig::default()); - let mut sink = RecordingSink::new(); - h.dispatch( - r#"{"type":"terminal_open","session_id":"s1","rows":24,"cols":80}"#, - &mut sink, - ) - .await - .unwrap(); - h.dispatch( - r#"{"type":"terminal_input","session_id":"s1","data":"aGk="}"#, - &mut sink, - ) - .await - .unwrap(); - h.dispatch( - r#"{"type":"terminal_resize","session_id":"s1","rows":30,"cols":100}"#, - &mut sink, - ) - .await - .unwrap(); - h.dispatch( - r#"{"type":"terminal_close","session_id":"s1"}"#, - &mut sink, - ) - .await - .unwrap(); - assert_eq!(sink.sent_count(), 0, "terminal control writes nothing to the WS sink"); - } - - #[tokio::test] - async fn test_dispatch_upgrade_denied_when_capability_absent() { - // CAP_UPGRADE off -> denied is written DIRECTLY to the sink (not the - // channel), unlike Exec. - let caps = ALL_CAPS & !serverbee_common::constants::CAP_UPGRADE; - let mut h = Harness::new(caps, FileConfig::default()); - let mut sink = RecordingSink::new(); - h.dispatch( - r#"{"type":"upgrade","version":"9.9.9","job_id":"j1"}"#, - &mut sink, - ) - .await - .unwrap(); - let msgs = sink.agent_messages(); - assert_eq!(msgs.len(), 1); - match &msgs[0] { - AgentMessage::CapabilityDenied { - capability, reason, .. - } => { - assert_eq!(capability, "upgrade"); - assert_eq!(*reason, CapabilityDeniedReason::AgentCapabilityDisabled); - } - other => panic!("expected CapabilityDenied, got {other:?}"), - } - } - - #[tokio::test] - async fn test_dispatch_traceroute_denied_when_capability_absent() { - let caps = ALL_CAPS & !serverbee_common::constants::CAP_PING_ICMP; - let mut h = Harness::new(caps, FileConfig::default()); - let mut sink = RecordingSink::new(); - h.dispatch( - r#"{"type":"traceroute","request_id":"r1","target":"example.com","max_hops":30}"#, - &mut sink, - ) - .await - .unwrap(); - assert_eq!(sink.sent_count(), 0, "denied traceroute goes to the channel"); - let denied = h.cmd_result_rx.recv().await.expect("denied expected"); - match denied { - AgentMessage::CapabilityDenied { - msg_id, - capability, - reason, - .. - } => { - assert_eq!(msg_id, Some("r1".to_string())); - assert_eq!(capability, "ping_icmp"); - assert_eq!(reason, CapabilityDeniedReason::AgentCapabilityDisabled); - } - other => panic!("expected CapabilityDenied, got {other:?}"), - } - } - - #[tokio::test] - async fn test_dispatch_traceroute_invalid_target_is_rejected() { - // Capability present but target fails validation -> a completed - // TracerouteRoundUpdate with an error is emitted on the channel. No - // real traceroute subprocess is spawned. - let mut h = Harness::new(ALL_CAPS, FileConfig::default()); - let mut sink = RecordingSink::new(); - h.dispatch( - r#"{"type":"traceroute","request_id":"r2","target":"bad target with spaces; rm -rf","max_hops":30}"#, - &mut sink, - ) - .await - .unwrap(); - assert_eq!(sink.sent_count(), 0); - let msg = h.cmd_result_rx.recv().await.expect("update expected"); - match msg { - AgentMessage::TracerouteRoundUpdate { - request_id, - completed, - error, - .. - } => { - assert_eq!(request_id, "r2"); - assert!(completed); - assert!(error.is_some(), "invalid target must carry an error"); - } - other => panic!("expected TracerouteRoundUpdate, got {other:?}"), - } - } - - #[tokio::test] - async fn test_dispatch_file_ops_denied_when_capability_absent() { - // CAP_FILE off -> each file op replies with a disabled error frame on - // the sink (capability-absent branch). - let caps = ALL_CAPS & !serverbee_common::constants::CAP_FILE; - let mut h = Harness::new(caps, FileConfig::default()); - let mut sink = RecordingSink::new(); - - h.dispatch(r#"{"type":"file_list","msg_id":"m1","path":"/tmp"}"#, &mut sink) - .await - .unwrap(); - h.dispatch(r#"{"type":"file_stat","msg_id":"m2","path":"/tmp"}"#, &mut sink) - .await - .unwrap(); - h.dispatch( - r#"{"type":"file_read","msg_id":"m3","path":"/tmp/x","max_size":1024}"#, - &mut sink, - ) - .await - .unwrap(); - h.dispatch( - r#"{"type":"file_write","msg_id":"m4","path":"/tmp/x","content":"aGk="}"#, - &mut sink, - ) - .await - .unwrap(); - h.dispatch( - r#"{"type":"file_delete","msg_id":"m5","path":"/tmp/x","recursive":false}"#, - &mut sink, - ) - .await - .unwrap(); - h.dispatch(r#"{"type":"file_mkdir","msg_id":"m6","path":"/tmp/d"}"#, &mut sink) - .await - .unwrap(); - h.dispatch( - r#"{"type":"file_move","msg_id":"m7","from":"/tmp/a","to":"/tmp/b"}"#, - &mut sink, - ) - .await - .unwrap(); - h.dispatch( - r#"{"type":"file_download_start","transfer_id":"t1","path":"/tmp/x"}"#, - &mut sink, - ) - .await - .unwrap(); - h.dispatch( - r#"{"type":"file_upload_start","transfer_id":"t2","path":"/tmp/x","size":4}"#, - &mut sink, - ) - .await - .unwrap(); - h.dispatch( - r#"{"type":"file_upload_chunk","transfer_id":"t3","offset":0,"data":"aGk="}"#, - &mut sink, - ) - .await - .unwrap(); - h.dispatch( - r#"{"type":"file_upload_end","transfer_id":"t4"}"#, - &mut sink, - ) - .await - .unwrap(); - - let msgs = sink.agent_messages(); - // 11 dispatches each produce exactly one response frame. - assert_eq!(msgs.len(), 11, "each denied file op emits one frame"); - // Spot-check representative variants carry the disabled error. - assert!(msgs.iter().any(|m| matches!( - m, - AgentMessage::FileListResult { error: Some(e), .. } if e.contains("disabled") - ))); - assert!(msgs.iter().any(|m| matches!( - m, - AgentMessage::FileOpResult { success: false, error: Some(e), .. } if e.contains("disabled") - ))); - assert!(msgs.iter().any(|m| matches!( - m, - AgentMessage::FileDownloadError { error, .. } if error.contains("disabled") - ))); - assert!(msgs.iter().any(|m| matches!( - m, - AgentMessage::FileUploadError { error, .. } if error.contains("disabled") - ))); - } - - #[tokio::test] - async fn test_dispatch_file_download_cancel_is_silent_noop() { - // FileDownloadCancel has no capability gate and no response; it just - // calls cancel_download. Cancelling an unknown transfer is a no-op. - let mut h = Harness::new(ALL_CAPS, FileConfig::default()); - let mut sink = RecordingSink::new(); - h.dispatch( - r#"{"type":"file_download_cancel","transfer_id":"nope"}"#, - &mut sink, - ) - .await - .unwrap(); - assert_eq!(sink.sent_count(), 0); - } - - #[tokio::test] - async fn test_dispatch_file_ops_success_with_enabled_manager() { - // File manager enabled with a real temp root: exercise the success - // branches (mkdir -> write -> list -> stat -> read -> move -> delete). - let tmp = tempfile::tempdir().unwrap(); - let root = std::fs::canonicalize(tmp.path()).unwrap(); - let cfg = enabled_file_cfg(&root); - let mut h = Harness::new(ALL_CAPS, cfg); - let mut sink = RecordingSink::new(); - - let sub = root.join("sub"); - let file_a = sub.join("a.txt"); - let file_b = sub.join("b.txt"); - let mkdir = format!( - r#"{{"type":"file_mkdir","msg_id":"mk","path":"{}"}}"#, - sub.to_string_lossy() - ); - h.dispatch(&mkdir, &mut sink).await.unwrap(); - - // `validate_path` canonicalizes, which requires the target to already - // exist — mirror the file_manager's own tests by pre-creating an empty - // file so the write overwrites it. - std::fs::write(&file_a, "").unwrap(); - - // base64("hi") == "aGk=" - let write = format!( - r#"{{"type":"file_write","msg_id":"w","path":"{}","content":"aGk="}}"#, - file_a.to_string_lossy() - ); - h.dispatch(&write, &mut sink).await.unwrap(); - - let list = format!( - r#"{{"type":"file_list","msg_id":"ls","path":"{}"}}"#, - sub.to_string_lossy() - ); - h.dispatch(&list, &mut sink).await.unwrap(); - - let stat = format!( - r#"{{"type":"file_stat","msg_id":"st","path":"{}"}}"#, - file_a.to_string_lossy() - ); - h.dispatch(&stat, &mut sink).await.unwrap(); - - let read = format!( - r#"{{"type":"file_read","msg_id":"rd","path":"{}","max_size":1024}}"#, - file_a.to_string_lossy() - ); - h.dispatch(&read, &mut sink).await.unwrap(); - - let mv = format!( - r#"{{"type":"file_move","msg_id":"mv","from":"{}","to":"{}"}}"#, - file_a.to_string_lossy(), - file_b.to_string_lossy() - ); - h.dispatch(&mv, &mut sink).await.unwrap(); - - let del = format!( - r#"{{"type":"file_delete","msg_id":"del","path":"{}","recursive":false}}"#, - file_b.to_string_lossy() - ); - h.dispatch(&del, &mut sink).await.unwrap(); - - let msgs = sink.agent_messages(); - assert_eq!(msgs.len(), 7, "seven file ops, seven responses"); - - // mkdir succeeded - assert!(msgs.iter().any(|m| matches!( - m, - AgentMessage::FileOpResult { msg_id, success: true, .. } if msg_id == "mk" - ))); - // write succeeded - assert!(msgs.iter().any(|m| matches!( - m, - AgentMessage::FileOpResult { msg_id, success: true, .. } if msg_id == "w" - ))); - // list returned at least the written file, no error - assert!(msgs.iter().any(|m| matches!( - m, - AgentMessage::FileListResult { msg_id, error: None, entries, .. } - if msg_id == "ls" && entries.iter().any(|e| e.name == "a.txt") - ))); - // stat found the entry - assert!(msgs.iter().any(|m| matches!( - m, - AgentMessage::FileStatResult { msg_id, entry: Some(_), error: None } if msg_id == "st" - ))); - // read returned the base64 content of "hi" - assert!(msgs.iter().any(|m| matches!( - m, - AgentMessage::FileReadResult { msg_id, content: Some(c), error: None } - if msg_id == "rd" && c == "aGk=" - ))); - // move + delete succeeded - assert!(msgs.iter().any(|m| matches!( - m, - AgentMessage::FileOpResult { msg_id, success: true, .. } if msg_id == "mv" - ))); - assert!(msgs.iter().any(|m| matches!( - m, - AgentMessage::FileOpResult { msg_id, success: true, .. } if msg_id == "del" - ))); - } - - #[tokio::test] - async fn test_dispatch_file_upload_success_round_trip() { - // Enabled manager: start -> chunk -> end upload, all on the sink. - let tmp = tempfile::tempdir().unwrap(); - let root = std::fs::canonicalize(tmp.path()).unwrap(); - let cfg = enabled_file_cfg(&root); - let mut h = Harness::new(ALL_CAPS, cfg); - let mut sink = RecordingSink::new(); - - let dest = root.join("up.bin"); - // "hi" -> base64 "aGk=" -> 2 bytes - let start = format!( - r#"{{"type":"file_upload_start","transfer_id":"u1","path":"{}","size":2}}"#, - dest.to_string_lossy() - ); - h.dispatch(&start, &mut sink).await.unwrap(); - h.dispatch( - r#"{"type":"file_upload_chunk","transfer_id":"u1","offset":0,"data":"aGk="}"#, - &mut sink, - ) - .await - .unwrap(); - h.dispatch( - r#"{"type":"file_upload_end","transfer_id":"u1"}"#, - &mut sink, - ) - .await - .unwrap(); - - let msgs = sink.agent_messages(); - assert_eq!(msgs.len(), 3); - // start ack at offset 0 - assert!(matches!( - &msgs[0], - AgentMessage::FileUploadAck { transfer_id, offset: 0 } if transfer_id == "u1" - )); - // chunk ack advances offset to 2 - assert!(matches!( - &msgs[1], - AgentMessage::FileUploadAck { transfer_id, offset: 2 } if transfer_id == "u1" - )); - // upload complete - assert!(matches!( - &msgs[2], - AgentMessage::FileUploadComplete { transfer_id } if transfer_id == "u1" - )); - // bytes actually landed on disk - assert_eq!(std::fs::read(&dest).unwrap(), b"hi"); - } - - #[tokio::test] - async fn test_dispatch_docker_start_stats_unavailable_emits_unavailable() { - // docker_manager is None -> DockerStartStats replies DockerUnavailable - // and leaves the stats interval unset. - let mut h = Harness::new(ALL_CAPS, FileConfig::default()); - let mut sink = RecordingSink::new(); - h.dispatch( - r#"{"type":"docker_start_stats","interval_secs":2}"#, - &mut sink, - ) - .await - .unwrap(); - let msgs = sink.agent_messages(); - assert_eq!(msgs.len(), 1); - assert!(matches!( - &msgs[0], - AgentMessage::DockerUnavailable { msg_id: None } - )); - assert!(h.docker_stats_interval.is_none()); - } - - #[tokio::test] - async fn test_dispatch_docker_stop_stats_clears_interval() { - let mut h = Harness::new(ALL_CAPS, FileConfig::default()); - // Pre-seed an interval so we can observe it being cleared. - h.docker_stats_interval = - Some(tokio::time::interval(Duration::from_secs(60))); - let mut sink = RecordingSink::new(); - h.dispatch(r#"{"type":"docker_stop_stats"}"#, &mut sink) - .await - .unwrap(); - assert_eq!(sink.sent_count(), 0); - assert!(h.docker_stats_interval.is_none()); - } - - #[tokio::test] - async fn test_dispatch_docker_request_unavailable_carries_msg_id() { - // Request variants with docker_manager None reply DockerUnavailable - // echoing the request's msg_id. - let mut h = Harness::new(ALL_CAPS, FileConfig::default()); - let mut sink = RecordingSink::new(); - h.dispatch( - r#"{"type":"docker_list_containers","msg_id":"req-1"}"#, - &mut sink, - ) - .await - .unwrap(); - let msgs = sink.agent_messages(); - assert_eq!(msgs.len(), 1); - assert!(matches!( - &msgs[0], - AgentMessage::DockerUnavailable { msg_id: Some(id) } if id == "req-1" - )); - - // An event variant with no msg_id replies with msg_id: None. - let mut sink2 = RecordingSink::new(); - h.dispatch(r#"{"type":"docker_events_start"}"#, &mut sink2) - .await - .unwrap(); - let msgs2 = sink2.agent_messages(); - assert_eq!(msgs2.len(), 1); - assert!(matches!( - &msgs2[0], - AgentMessage::DockerUnavailable { msg_id: None } - )); - } - - #[tokio::test] - async fn test_dispatch_blocklist_reset_returns_ack() { - // FirewallManager uses the real CliNftExecutor. On a host without `nft` - // (macOS CI) BlocklistReset deterministically fails the wipe but still - // returns a BlocklistResetAck reply, which the dispatcher forwards. - let mut h = Harness::new(ALL_CAPS, FileConfig::default()); - let mut sink = RecordingSink::new(); - h.dispatch(r#"{"type":"blocklist_reset"}"#, &mut sink) - .await - .unwrap(); - let msgs = sink.agent_messages(); - assert_eq!(msgs.len(), 1, "reset always produces an ack"); - assert!(matches!( - &msgs[0], - AgentMessage::BlocklistResetAck { .. } - )); - } - - #[tokio::test] - async fn test_dispatch_ip_quality_sync_and_run_now_respect_capability() { - // With CAP_IP_QUALITY present, sync/run_now are accepted (no WS output). - let mut h = Harness::new(ALL_CAPS, FileConfig::default()); - let mut sink = RecordingSink::new(); - h.dispatch( - r#"{"type":"ip_quality_sync","services":[],"interval_hours":12}"#, - &mut sink, - ) - .await - .unwrap(); - h.dispatch(r#"{"type":"ip_quality_run_now"}"#, &mut sink) - .await - .unwrap(); - assert_eq!(sink.sent_count(), 0); - - // Without CAP_IP_QUALITY, both are silently ignored as well. - let caps = ALL_CAPS & !serverbee_common::constants::CAP_IP_QUALITY; - let mut h2 = Harness::new(caps, FileConfig::default()); - let mut sink2 = RecordingSink::new(); - h2.dispatch( - r#"{"type":"ip_quality_sync","services":[],"interval_hours":6}"#, - &mut sink2, - ) - .await - .unwrap(); - h2.dispatch(r#"{"type":"ip_quality_run_now"}"#, &mut sink2) - .await - .unwrap(); - assert_eq!(sink2.sent_count(), 0); - } - - #[tokio::test] - async fn test_dispatch_blocklist_sync_add_remove_forward_acks() { - // Sync/Add/Remove all route into FirewallManager and forward its - // BlocklistAck reply over the WS sink. On a host without `nft` the - // apply fails but the manager still returns an ack (with Failed state), - // so the dispatcher always emits exactly one frame per request. - let mut h = Harness::new(ALL_CAPS, FileConfig::default()); - - // Full-state sync with one entry. - let mut sink = RecordingSink::new(); - h.dispatch( - r#"{"type":"blocklist_sync","entries":[{"id":"e1","target":"1.2.3.4/32","family":4}]}"#, - &mut sink, - ) - .await - .unwrap(); - let msgs = sink.agent_messages(); - assert_eq!(msgs.len(), 1, "sync emits one ack frame"); - assert!(matches!(&msgs[0], AgentMessage::BlocklistAck { .. })); - - // Incremental add. - let mut sink = RecordingSink::new(); - h.dispatch( - r#"{"type":"blocklist_add","entry":{"id":"e2","target":"5.6.7.8/32","family":4}}"#, - &mut sink, - ) - .await - .unwrap(); - let msgs = sink.agent_messages(); - assert_eq!(msgs.len(), 1, "add emits one ack frame"); - assert!(matches!(&msgs[0], AgentMessage::BlocklistAck { .. })); - - // Incremental remove of an unknown id still produces a single-item ack. - let mut sink = RecordingSink::new(); - h.dispatch(r#"{"type":"blocklist_remove","id":"e2"}"#, &mut sink) - .await - .unwrap(); - let msgs = sink.agent_messages(); - assert_eq!(msgs.len(), 1, "remove emits one ack frame"); - assert!(matches!(&msgs[0], AgentMessage::BlocklistAck { .. })); - } - - #[tokio::test] - async fn test_dispatch_file_ops_disabled_manager_replies_disabled_even_with_capability() { - // CAP_FILE present but the manager is disabled (default FileConfig has - // enabled=false). The `!file_manager.is_enabled()` half of the guard - // must still short-circuit with the disabled error, independent of caps. - let mut h = Harness::new(ALL_CAPS, FileConfig::default()); - assert!(!h.file_manager.is_enabled(), "default file manager is disabled"); - let mut sink = RecordingSink::new(); - h.dispatch(r#"{"type":"file_list","msg_id":"m1","path":"/tmp"}"#, &mut sink) - .await - .unwrap(); - h.dispatch( - r#"{"type":"file_write","msg_id":"m2","path":"/tmp/x","content":"aGk="}"#, - &mut sink, - ) - .await - .unwrap(); - let msgs = sink.agent_messages(); - assert_eq!(msgs.len(), 2, "each op replies once even with cap present"); - assert!(msgs.iter().any(|m| matches!( - m, - AgentMessage::FileListResult { error: Some(e), .. } if e.contains("disabled") - ))); - assert!(msgs.iter().any(|m| matches!( - m, - AgentMessage::FileOpResult { success: false, error: Some(e), .. } if e.contains("disabled") - ))); - } - - #[tokio::test] - async fn test_dispatch_upgrade_already_running_emits_failure_on_channel() { - // Force the global single-flight latch to "in progress", then dispatch - // an Upgrade with the capability present. The duplicate must be rejected - // with an UpgradeResult error on the cmd channel (not the WS sink), and - // the latch must be left untouched (still true) for the real holder. - UPGRADE_IN_PROGRESS.store(true, Ordering::SeqCst); - // Ensure we always release the global latch so other tests aren't poisoned. - struct Guard; - impl Drop for Guard { - fn drop(&mut self) { - UPGRADE_IN_PROGRESS.store(false, Ordering::SeqCst); - } - } - let _guard = Guard; - - let mut h = Harness::new(ALL_CAPS, FileConfig::default()); - let mut sink = RecordingSink::new(); - h.dispatch( - r#"{"type":"upgrade","version":"9.9.9","job_id":"dup-job"}"#, - &mut sink, - ) - .await - .unwrap(); - // Nothing is written to the WS sink for the duplicate case. - assert_eq!(sink.sent_count(), 0, "duplicate upgrade writes to channel, not sink"); - let msg = tokio::time::timeout(Duration::from_secs(5), h.cmd_result_rx.recv()) - .await - .expect("failure msg expected in time") - .expect("UpgradeResult expected"); - match msg { - AgentMessage::UpgradeResult { - job_id, - target_version, - stage, - error, - .. - } => { - assert_eq!(job_id, Some("dup-job".to_string())); - assert_eq!(target_version, "9.9.9"); - assert_eq!(stage, UpgradeStage::Downloading); - assert!(error.contains("already running")); - } - other => panic!("expected UpgradeResult, got {other:?}"), - } - } - - // ---------------------------------------------------------------------- - // `docker_request_msg_id` — remaining non-msg-id docker variants. - // ---------------------------------------------------------------------- - - #[test] - fn test_docker_request_msg_id_none_for_log_and_event_variants() { - // Streaming control variants carry no request msg_id -> None. - assert_eq!( - docker_request_msg_id(&ServerMessage::DockerLogsStart { - session_id: "s".to_string(), - container_id: "c".to_string(), - tail: None, - follow: false, - }), - None - ); - assert_eq!( - docker_request_msg_id(&ServerMessage::DockerLogsStop { - session_id: "s".to_string(), - }), - None - ); - assert_eq!(docker_request_msg_id(&ServerMessage::DockerEventsStart), None); - assert_eq!(docker_request_msg_id(&ServerMessage::DockerEventsStop), None); - assert_eq!(docker_request_msg_id(&ServerMessage::DockerStartStats { interval_secs: 5 }), None); - } - - // ---------------------------------------------------------------------- - // `execute_command` — pure-ish process helper. These shell out to `sh`, - // which is always present on macOS/Linux CI, so they remain deterministic. - // ---------------------------------------------------------------------- - - #[tokio::test] - async fn test_execute_command_captures_stdout_and_zero_exit() { - // Deterministic stdout, exit 0. - let r = execute_command("t-ok", "printf hello", Some(5)).await; - assert_eq!(r.task_id, "t-ok"); - assert_eq!(r.exit_code, 0); - assert!(r.output.contains("hello")); - } - - #[tokio::test] - async fn test_execute_command_nonzero_exit_and_stderr_appended() { - // `sh -c 'exit 3'` yields exit_code 3; stderr is folded into output. - let r = execute_command("t-fail", "echo oops 1>&2; exit 3", Some(5)).await; - assert_eq!(r.exit_code, 3); - assert!(r.output.contains("oops"), "stderr must be appended to output"); - } - - #[tokio::test] - async fn test_execute_command_truncates_large_output() { - // Emit more than MAX_TASK_OUTPUT_SIZE bytes; the helper must cap and - // append the truncation marker. `yes | head -c N` is portable. - let cmd = format!("yes A | head -c {}", MAX_TASK_OUTPUT_SIZE + 5000); - let r = execute_command("t-big", &cmd, Some(10)).await; - assert_eq!(r.exit_code, 0); - assert!( - r.output.ends_with("\n... (output truncated)"), - "oversized output must carry the truncation marker" - ); - assert!( - r.output.len() <= MAX_TASK_OUTPUT_SIZE + "\n... (output truncated)".len(), - "truncated output must respect the cap" - ); - } - - #[tokio::test] - async fn test_execute_command_times_out_with_negative_exit() { - // A 2s sleep against a 1s timeout must surface the timeout branch. - let r = execute_command("t-timeout", "sleep 2", Some(1)).await; - assert_eq!(r.exit_code, -1); - assert!(r.output.contains("timed out"), "timeout branch must report it"); - } - // ---------------------------------------------------------------------- // Upgrade progress/failure emitters — pure channel-send helpers. // ---------------------------------------------------------------------- @@ -2702,6 +1106,17 @@ mod tests { /// Server-side half of an accepted fake connection. type ServerWs = WebSocketStream; + /// All capability bits set — every success arm runs. + const ALL_CAPS: u32 = serverbee_common::constants::CAP_VALID_MASK; + + fn enabled_file_cfg(root: &std::path::Path) -> FileConfig { + FileConfig { + enabled: true, + root_paths: vec![root.to_string_lossy().to_string()], + ..FileConfig::default() + } + } + /// Build a reporter config that points at the given loopback `ws_addr` /// (`host:port`) and disables every network-touching background task so /// the fake server only ever sees frames produced by the connect/report @@ -3461,9 +1876,10 @@ mod tests { }, ) .await; - let reply = - read_agent_until(&mut ws, |m| matches!(m, AgentMessage::FileListResult { .. })) - .await; + let reply = read_agent_until(&mut ws, |m| { + matches!(m, AgentMessage::FileListResult { .. }) + }) + .await; ws.send(WsMessage::Close(None)).await.ok(); reply }); diff --git a/crates/agent/src/reporter/runtime.rs b/crates/agent/src/reporter/runtime.rs new file mode 100644 index 000000000..a7be69d85 --- /dev/null +++ b/crates/agent/src/reporter/runtime.rs @@ -0,0 +1,1701 @@ +//! Connection-scoped command runtime. +//! +//! One WebSocket connection owns one [`ConnectionRuntime`]: every manager, +//! outbound channel sender, and piece of Docker lifecycle state that executing +//! a `ServerMessage` can touch lives here, so the dispatcher is a method on +//! the state it drives instead of a free function threading a dozen borrows. +//! `mod.rs` keeps transport concerns only (connect/backoff, the select! loop, +//! report and IP timers); file replies stay in `file_ops`, framing in `wire`. + +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; + +use futures_util::SinkExt; +use serverbee_common::constants::{DEFAULT_COMMAND_TIMEOUT_SECS, MAX_TASK_OUTPUT_SIZE}; +use serverbee_common::protocol::{AgentMessage, ServerMessage, UpgradeStage}; +use serverbee_common::types::{NetworkProbeResultData, PingResult}; +use tokio::sync::mpsc; +use tokio_tungstenite::tungstenite::Message; + +use super::wire::send_msg; +use super::{emit_upgrade_failure, file_ops, perform_upgrade}; +use crate::capability_grants::CapabilityAuthority; +use crate::config::{AgentConfig, UpgradeConfig}; +use crate::docker::DockerManager; +use crate::file_manager::{FileEvent, FileManager}; +use crate::firewall::FirewallManager; +use crate::ip_quality::{RunResult, UnlockChecker}; +use crate::network_prober::NetworkProber; +use crate::pinger::PingManager; +use crate::terminal::{TerminalEvent, TerminalManager}; + +const DOCKER_RETRY_SECS: u64 = 30; + +/// Process-wide single-flight latch for binary self-upgrade. A duplicate +/// Upgrade command must be rejected without disturbing the running one. +static UPGRADE_IN_PROGRESS: AtomicBool = AtomicBool::new(false); + +/// Which Docker housekeeping event [`ConnectionRuntime::docker_tick`] resolved +/// to: poll container stats, or retry the daemon connection. +pub(super) enum DockerTick { + PollStats, + Retry, +} + +/// Receiver halves for every event stream the runtime's managers emit. The +/// select! loop in `mod.rs` drains these; the matching senders live inside +/// [`ConnectionRuntime`] so a handler can never outlive its channel. +pub(super) struct ConnectionEvents { + pub(super) ping_rx: mpsc::Receiver, + pub(super) term_rx: mpsc::Receiver, + pub(super) network_probe_rx: mpsc::Receiver, + pub(super) file_rx: mpsc::Receiver, + pub(super) unlock_result_rx: mpsc::Receiver, + pub(super) docker_rx: mpsc::Receiver, + pub(super) cmd_result_rx: mpsc::Receiver, +} + +/// Everything a single agent connection needs to execute server commands. +pub(super) struct ConnectionRuntime { + pub(super) capabilities: Arc, + pub(super) ping_manager: PingManager, + pub(super) terminal_manager: TerminalManager, + pub(super) network_prober: NetworkProber, + pub(super) file_manager: FileManager, + pub(super) unlock_checker: UnlockChecker, + pub(super) cmd_result_tx: mpsc::Sender, + file_tx: mpsc::Sender, + docker_tx: mpsc::Sender, + pub(super) docker_manager: Option, + docker_available: bool, + pub(super) docker_stats_interval: Option, + docker_retry_interval: tokio::time::Interval, + upgrade_cfg: UpgradeConfig, + firewall_manager: Arc, +} + +impl ConnectionRuntime { + /// Build the runtime and hand back the event receivers the select! loop + /// drains. Docker starts absent; call [`Self::probe_docker`] before + /// advertising features. + pub(super) fn new( + config: &AgentConfig, + capabilities: Arc, + firewall_manager: Arc, + ) -> (Self, ConnectionEvents) { + let (ping_tx, ping_rx) = mpsc::channel(256); + let ping_manager = PingManager::new(ping_tx, Arc::clone(&capabilities)); + + let (term_tx, term_rx) = mpsc::channel(256); + let terminal_manager = TerminalManager::new(term_tx, Arc::clone(&capabilities)); + + let (network_probe_tx, network_probe_rx) = mpsc::channel::(256); + let network_prober = NetworkProber::new(network_probe_tx, Arc::clone(&capabilities)); + + let (unlock_result_tx, unlock_result_rx) = mpsc::channel::(8); + let unlock_checker = UnlockChecker::new(Arc::clone(&capabilities), unlock_result_tx); + + let (file_tx, file_rx) = mpsc::channel::(16); + let file_manager = FileManager::new(config.file.clone(), Arc::clone(&capabilities)); + + let (docker_tx, docker_rx) = mpsc::channel::(256); + + let (cmd_result_tx, cmd_result_rx) = mpsc::channel::(32); + + // First retry fires one full period out — connection start already + // probes explicitly, so there is no immediate tick to consume. + let retry_period = Duration::from_secs(DOCKER_RETRY_SECS); + let docker_retry_interval = + tokio::time::interval_at(tokio::time::Instant::now() + retry_period, retry_period); + + let runtime = Self { + capabilities, + ping_manager, + terminal_manager, + network_prober, + file_manager, + unlock_checker, + cmd_result_tx, + file_tx, + docker_tx, + docker_manager: None, + docker_available: false, + docker_stats_interval: None, + docker_retry_interval, + upgrade_cfg: config.upgrade.clone(), + firewall_manager, + }; + let events = ConnectionEvents { + ping_rx, + term_rx, + network_probe_rx, + file_rx, + unlock_result_rx, + docker_rx, + cmd_result_rx, + }; + (runtime, events) + } + + /// Feature list advertised in `SystemInfo`, derived from what is live. + pub(super) fn features(&self) -> Vec { + if self.docker_available { + vec!["docker".to_string()] + } else { + Vec::new() + } + } + + /// Initial Docker detection at connection start. Absence is informational + /// — the retry tick keeps looking. + pub(super) async fn probe_docker(&mut self) { + match DockerManager::try_new(self.docker_tx.clone(), Arc::clone(&self.capabilities)) { + Ok(dm) => match dm.verify_connection().await { + Ok(()) => { + tracing::info!("Docker daemon connected"); + self.docker_available = true; + self.docker_manager = Some(dm); + } + Err(e) => { + tracing::info!("Docker daemon not available: {e}"); + } + }, + Err(e) => { + tracing::info!("Docker not available: {e}"); + } + } + } + + /// Wait for the next Docker housekeeping event: a stats-poll tick while + /// the daemon is connected, or a retry tick while it is absent. Pending + /// when connected with stats polling off. + pub(super) async fn docker_tick(&mut self) -> DockerTick { + if self.docker_manager.is_some() { + match self.docker_stats_interval.as_mut() { + Some(iv) => { + iv.tick().await; + DockerTick::PollStats + } + None => std::future::pending().await, + } + } else { + self.docker_retry_interval.tick().await; + DockerTick::Retry + } + } + + /// Poll container stats once; a failing daemon demotes the runtime. + pub(super) async fn poll_docker_stats(&mut self, write: &mut S) -> anyhow::Result<()> + where + S: SinkExt + Unpin, + { + if let Some(dm) = self.docker_manager.as_mut() + && let Err(e) = dm.poll_stats().await + { + tracing::warn!("Docker stats polling failed: {e}"); + self.demote_docker(write).await?; + } + Ok(()) + } + + /// Retry the daemon connection; success re-advertises the feature. + pub(super) async fn retry_docker(&mut self, write: &mut S) -> anyhow::Result<()> + where + S: SinkExt + Unpin, + { + tracing::debug!("Retrying Docker connection..."); + match DockerManager::try_new(self.docker_tx.clone(), Arc::clone(&self.capabilities)) { + Ok(dm) => match dm.verify_connection().await { + Ok(()) => { + tracing::info!("Docker daemon now available"); + self.docker_manager = Some(dm); + self.docker_available = true; + let msg = AgentMessage::FeaturesUpdate { + features: vec!["docker".to_string()], + }; + send_msg(write, &msg).await?; + } + Err(e) => { + tracing::debug!("Docker still not available: {e}"); + } + }, + Err(e) => { + tracing::debug!("Docker still not available: {e}"); + } + } + Ok(()) + } + + async fn demote_docker(&mut self, write: &mut S) -> anyhow::Result<()> + where + S: SinkExt + Unpin, + { + if let Some(dm) = self.docker_manager.as_mut() { + dm.cleanup(); + } + self.docker_manager = None; + self.docker_stats_interval = None; + + if self.docker_available { + self.docker_available = false; + let msg = AgentMessage::FeaturesUpdate { features: vec![] }; + send_msg(write, &msg).await?; + } + + Ok(()) + } + + /// Tear down every in-flight resource this connection owns. Called on all + /// connection-exit paths (server close, WS error, stream end). + pub(super) fn shutdown(&mut self) { + self.ping_manager.stop_all(); + self.terminal_manager.close_all(); + self.network_prober.stop_all(); + self.unlock_checker.stop(); + self.file_manager.cancel_all_transfers(); + if let Some(dm) = self.docker_manager.as_mut() { + dm.cleanup(); + } + } + + /// Parse and execute one server text frame against this connection's + /// state. Unparseable frames are logged and swallowed — a protocol + /// mismatch must never tear the connection down. + pub(super) async fn handle_server_message( + &mut self, + text: &str, + write: &mut S, + ) -> anyhow::Result<()> + where + S: SinkExt + Unpin, + { + use serverbee_common::constants::*; + + let msg: ServerMessage = match serde_json::from_str(text) { + Ok(m) => m, + Err(e) => { + tracing::warn!("Failed to parse server message: {e}"); + return Ok(()); + } + }; + + match msg { + ServerMessage::Ping => { + send_msg(write, &AgentMessage::Pong).await?; + tracing::debug!("Responded to Ping with Pong"); + } + ServerMessage::Exec { + task_id, + command, + timeout, + } => { + if !self.capabilities.has(CAP_EXEC) { + tracing::warn!("Exec denied: capability disabled (task_id={task_id})"); + spawn_capability_denied(&self.cmd_result_tx, Some(task_id), "exec"); + return Ok(()); + } + tracing::info!("Executing command (task_id={task_id}): {command}"); + let tx = self.cmd_result_tx.clone(); + tokio::spawn(async move { + let result = execute_command(&task_id, &command, timeout).await; + let msg = AgentMessage::TaskResult { + msg_id: uuid::Uuid::new_v4().to_string(), + result, + }; + if tx.send(msg).await.is_err() { + tracing::warn!( + "Failed to send TaskResult for task_id={task_id}: channel closed" + ); + } else { + tracing::info!("TaskResult ready for task_id={task_id}"); + } + }); + } + ServerMessage::Ack { msg_id } => { + tracing::debug!("Received Ack for msg_id={msg_id}"); + } + ServerMessage::Welcome { .. } => { + tracing::warn!("Unexpected second Welcome message"); + } + ServerMessage::PingTasksSync { tasks } => { + tracing::info!("Received PingTasksSync with {} tasks", tasks.len()); + self.ping_manager.sync(tasks); + } + ServerMessage::TerminalOpen { + session_id, + rows, + cols, + } => { + tracing::info!("Opening terminal session {session_id} ({cols}x{rows})"); + self.terminal_manager.open(session_id, rows, cols); + } + ServerMessage::TerminalInput { session_id, data } => { + self.terminal_manager.write_input(&session_id, &data); + } + ServerMessage::TerminalResize { + session_id, + rows, + cols, + } => { + tracing::debug!("Resizing terminal {session_id} to {cols}x{rows}"); + self.terminal_manager.resize(&session_id, rows, cols); + } + ServerMessage::TerminalClose { session_id } => { + tracing::debug!("Closing terminal session {session_id}"); + self.terminal_manager.close(&session_id); + } + ServerMessage::Upgrade { + version, job_id, .. + } => { + if !self.capabilities.has(CAP_UPGRADE) { + tracing::warn!("Upgrade denied: capability disabled"); + send_msg(write, &capability_denied_msg(None, "upgrade")).await?; + return Ok(()); + } + + if UPGRADE_IN_PROGRESS + .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) + .is_err() + { + let tx = self.cmd_result_tx.clone(); + tokio::spawn(async move { + emit_upgrade_failure( + &tx, + job_id, + version, + UpgradeStage::Downloading, + "another upgrade is already running".to_string(), + None, + ) + .await; + }); + return Ok(()); + } + + tracing::info!("Upgrade requested: v{version} (pinned source)"); + let upgrade_cfg = self.upgrade_cfg.clone(); + let tx = self.cmd_result_tx.clone(); + tokio::spawn(async move { + if let Err(e) = + perform_upgrade(&version, &upgrade_cfg, job_id, tx.clone()).await + { + tracing::error!("Upgrade to v{version} failed: {e}"); + UPGRADE_IN_PROGRESS.store(false, Ordering::SeqCst); + } + }); + } + ServerMessage::NetworkProbeSync { + targets, + interval, + packet_count, + } => { + tracing::info!( + "Received NetworkProbeSync: {} targets, interval={}s, packet_count={}", + targets.len(), + interval, + packet_count + ); + self.network_prober.sync(targets, interval, packet_count); + } + ServerMessage::Traceroute { + request_id, + target, + max_hops, + protocol, + } => { + if !self.capabilities.has(CAP_PING_ICMP) { + tracing::warn!( + "Traceroute denied: capability disabled (request_id={request_id})" + ); + spawn_capability_denied(&self.cmd_result_tx, Some(request_id), "ping_icmp"); + return Ok(()); + } + + // Input validation: target must be domain or IP only. + if !crate::traceroute::is_valid_traceroute_target(&target) { + tracing::warn!( + "Traceroute rejected: invalid target '{target}' (request_id={request_id})" + ); + let tx = self.cmd_result_tx.clone(); + let request_id_c = request_id.clone(); + let target_c = target.clone(); + tokio::spawn(async move { + let _ = tx + .send(AgentMessage::TracerouteRoundUpdate { + request_id: request_id_c, + target: target_c, + round: 0, + total_rounds: 0, + hops: vec![], + completed: true, + error: Some( + "Invalid target: must be a domain or IP address".into(), + ), + }) + .await; + }); + return Ok(()); + } + + let proto = protocol.unwrap_or(serverbee_common::protocol::TraceProtocol::Icmp); + tracing::info!( + "Executing traceroute to {target} (max_hops={max_hops}, request_id={request_id}, protocol={proto:?})" + ); + crate::traceroute::spawn_traceroute( + request_id, + target, + max_hops, + proto, + self.cmd_result_tx.clone(), + ); + } + // --- File management messages --- (gate + reply shapes live in file_ops) + msg @ (ServerMessage::FileList { .. } + | ServerMessage::FileStat { .. } + | ServerMessage::FileRead { .. } + | ServerMessage::FileWrite { .. } + | ServerMessage::FileDelete { .. } + | ServerMessage::FileMkdir { .. } + | ServerMessage::FileMove { .. } + | ServerMessage::FileDownloadStart { .. } + | ServerMessage::FileDownloadCancel { .. } + | ServerMessage::FileUploadStart { .. } + | ServerMessage::FileUploadChunk { .. } + | ServerMessage::FileUploadEnd { .. }) => { + file_ops::handle_file_message( + msg, + write, + &self.file_manager, + &self.file_tx, + &self.capabilities, + ) + .await?; + } + // --- Docker messages --- + ServerMessage::DockerStartStats { interval_secs } => { + if self.docker_manager.is_some() { + let secs = interval_secs.max(1); + tracing::info!("Starting Docker stats polling every {secs}s"); + let mut iv = tokio::time::interval(Duration::from_secs(secs as u64)); + iv.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + self.docker_stats_interval = Some(iv); + } else { + tracing::warn!("DockerStartStats received but Docker is not available"); + let unavailable = AgentMessage::DockerUnavailable { msg_id: None }; + send_msg(write, &unavailable).await?; + } + } + ServerMessage::DockerStopStats => { + tracing::info!("Stopping Docker stats polling"); + self.docker_stats_interval = None; + } + ServerMessage::DockerListContainers { .. } + | ServerMessage::DockerLogsStart { .. } + | ServerMessage::DockerLogsStop { .. } + | ServerMessage::DockerEventsStart + | ServerMessage::DockerEventsStop + | ServerMessage::DockerContainerAction { .. } + | ServerMessage::DockerGetInfo { .. } + | ServerMessage::DockerListNetworks { .. } + | ServerMessage::DockerListVolumes { .. } => { + if let Some(dm) = self.docker_manager.as_mut() { + if let Err(e) = dm.handle_server_message(msg.clone()).await { + tracing::warn!("Docker runtime became unavailable: {e}"); + self.demote_docker(write).await?; + } + } else { + tracing::warn!("Docker message received but Docker is not available"); + let unavailable = AgentMessage::DockerUnavailable { + msg_id: docker_request_msg_id(&msg), + }; + send_msg(write, &unavailable).await?; + } + } + // Firewall blocklist variants — dispatched to the FirewallManager + // state machine; any returned ack is sent straight back over the + // WebSocket. + // + // The mutating variants (Sync/Add/Remove) enforce CAP_FIREWALL_BLOCK + // on the agent's own host, mirroring the capability gates on Exec / + // File / Traceroute etc. — the server is not the only trust boundary. + // BlocklistReset is deliberately *not* gated: it wipes ServerBee's + // own nft table (cleanup / disable path) and must stay reachable even + // after the capability is revoked, so a denied agent can still be + // cleaned up. + msg @ (ServerMessage::BlocklistSync { .. } + | ServerMessage::BlocklistAdd { .. } + | ServerMessage::BlocklistRemove { .. } + | ServerMessage::BlocklistReset) => { + let is_reset = matches!(msg, ServerMessage::BlocklistReset); + if !is_reset && !self.capabilities.has(CAP_FIREWALL_BLOCK) { + tracing::warn!( + "Firewall blocklist mutation denied: CAP_FIREWALL_BLOCK not effective — ignoring" + ); + } else if let Some(reply) = self.firewall_manager.handle(msg).await { + send_msg(write, &reply).await?; + tracing::debug!("Sent firewall blocklist ack"); + } + } + ServerMessage::IpQualitySync { + services, + interval_hours, + } => { + if self.capabilities.has(CAP_IP_QUALITY) { + tracing::info!( + "Received IpQualitySync: {} services, interval={}h", + services.len(), + interval_hours + ); + self.unlock_checker.sync(services, interval_hours).await; + } else { + tracing::debug!( + "IpQualitySync received but CAP_IP_QUALITY not effective — ignoring" + ); + } + } + ServerMessage::IpQualityRunNow => { + if self.capabilities.has(CAP_IP_QUALITY) { + tracing::info!("Received IpQualityRunNow"); + self.unlock_checker.run_now(); + } else { + tracing::debug!( + "IpQualityRunNow received but CAP_IP_QUALITY not effective — ignoring" + ); + } + } + } + + Ok(()) + } +} + +/// Build the standard agent-local capability denial message. +fn capability_denied_msg(msg_id: Option, capability: &str) -> AgentMessage { + AgentMessage::CapabilityDenied { + msg_id, + session_id: None, + capability: capability.to_string(), + reason: serverbee_common::constants::CapabilityDeniedReason::AgentCapabilityDisabled, + } +} + +/// Queue a capability denial onto the outbound channel without blocking the +/// read loop. +fn spawn_capability_denied( + tx: &mpsc::Sender, + msg_id: Option, + capability: &str, +) { + let denied = capability_denied_msg(msg_id, capability); + let tx = tx.clone(); + tokio::spawn(async move { + let _ = tx.send(denied).await; + }); +} + +fn docker_request_msg_id(msg: &ServerMessage) -> Option { + match msg { + ServerMessage::DockerListContainers { msg_id } + | ServerMessage::DockerGetInfo { msg_id } + | ServerMessage::DockerListNetworks { msg_id } + | ServerMessage::DockerListVolumes { msg_id } => Some(msg_id.clone()), + ServerMessage::DockerContainerAction { msg_id, .. } => Some(msg_id.clone()), + _ => None, + } +} + +async fn execute_command( + task_id: &str, + command: &str, + timeout: Option, +) -> serverbee_common::types::TaskResult { + let timeout_secs = timeout.unwrap_or(DEFAULT_COMMAND_TIMEOUT_SECS); + + let result = tokio::time::timeout( + Duration::from_secs(timeout_secs as u64), + tokio::process::Command::new("sh") + .arg("-c") + .arg(command) + .output(), + ) + .await; + + match result { + Ok(Ok(output)) => { + let mut combined = String::from_utf8_lossy(&output.stdout).to_string(); + let stderr = String::from_utf8_lossy(&output.stderr); + if !stderr.is_empty() { + combined.push('\n'); + combined.push_str(&stderr); + } + if combined.len() > MAX_TASK_OUTPUT_SIZE { + combined.truncate(MAX_TASK_OUTPUT_SIZE); + combined.push_str("\n... (output truncated)"); + } + serverbee_common::types::TaskResult { + task_id: task_id.to_string(), + output: combined, + exit_code: output.status.code().unwrap_or(-1), + } + } + Ok(Err(e)) => serverbee_common::types::TaskResult { + task_id: task_id.to_string(), + output: format!("Failed to execute command: {e}"), + exit_code: -1, + }, + Err(_) => serverbee_common::types::TaskResult { + task_id: task_id.to_string(), + output: format!("Command timed out after {timeout_secs}s"), + exit_code: -1, + }, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::{ + CapabilitiesConfig, CollectorConfig, FileConfig, IpChangeConfig, LogConfig, SecurityConfig, + UpgradeConfig, + }; + use crate::firewall::nft::CliNftExecutor; + use serverbee_common::constants::CapabilityDeniedReason; + + #[test] + fn test_docker_request_msg_id_extracts_only_request_variants() { + // Variants that carry a msg_id return it... + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerListContainers { + msg_id: "a".to_string() + }), + Some("a".to_string()) + ); + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerGetInfo { + msg_id: "b".to_string() + }), + Some("b".to_string()) + ); + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerListNetworks { + msg_id: "c".to_string() + }), + Some("c".to_string()) + ); + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerListVolumes { + msg_id: "d".to_string() + }), + Some("d".to_string()) + ); + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerContainerAction { + msg_id: "e".to_string(), + container_id: "cid".to_string(), + action: serverbee_common::docker_types::DockerAction::Restart { timeout: None }, + }), + Some("e".to_string()) + ); + // ...non-request docker variants return None. + assert_eq!(docker_request_msg_id(&ServerMessage::DockerStopStats), None); + assert_eq!(docker_request_msg_id(&ServerMessage::Ping), None); + } + + #[test] + fn test_docker_request_msg_id_none_for_log_and_event_variants() { + // Streaming control variants carry no request msg_id -> None. + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerLogsStart { + session_id: "s".to_string(), + container_id: "c".to_string(), + tail: None, + follow: false, + }), + None + ); + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerLogsStop { + session_id: "s".to_string(), + }), + None + ); + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerEventsStart), + None + ); + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerEventsStop), + None + ); + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerStartStats { interval_secs: 5 }), + None + ); + } + + // ---------------------------------------------------------------------- + // `execute_command` — pure-ish process helper. These shell out to `sh`, + // which is always present on macOS/Linux CI, so they remain deterministic. + // ---------------------------------------------------------------------- + + #[tokio::test] + async fn test_execute_command_captures_stdout_and_zero_exit() { + // Deterministic stdout, exit 0. + let r = execute_command("t-ok", "printf hello", Some(5)).await; + assert_eq!(r.task_id, "t-ok"); + assert_eq!(r.exit_code, 0); + assert!(r.output.contains("hello")); + } + + #[tokio::test] + async fn test_execute_command_nonzero_exit_and_stderr_appended() { + // `sh -c 'exit 3'` yields exit_code 3; stderr is folded into output. + let r = execute_command("t-fail", "echo oops 1>&2; exit 3", Some(5)).await; + assert_eq!(r.exit_code, 3); + assert!( + r.output.contains("oops"), + "stderr must be appended to output" + ); + } + + #[tokio::test] + async fn test_execute_command_truncates_large_output() { + // Emit more than MAX_TASK_OUTPUT_SIZE bytes; the helper must cap and + // append the truncation marker. `yes | head -c N` is portable. + let cmd = format!("yes A | head -c {}", MAX_TASK_OUTPUT_SIZE + 5000); + let r = execute_command("t-big", &cmd, Some(10)).await; + assert_eq!(r.exit_code, 0); + assert!( + r.output.ends_with("\n... (output truncated)"), + "oversized output must carry the truncation marker" + ); + assert!( + r.output.len() <= MAX_TASK_OUTPUT_SIZE + "\n... (output truncated)".len(), + "truncated output must respect the cap" + ); + } + + #[tokio::test] + async fn test_execute_command_times_out_with_negative_exit() { + // A 2s sleep against a 1s timeout must surface the timeout branch. + let r = execute_command("t-timeout", "sleep 2", Some(1)).await; + assert_eq!(r.exit_code, -1); + assert!( + r.output.contains("timed out"), + "timeout branch must report it" + ); + } + + // ---------------------------------------------------------------------- + // `handle_server_message` dispatcher coverage via a mock sink. + // + // The dispatcher is generic over the WS sink (`S: SinkExt`), + // so we drive it with an in-memory recording sink instead of a real + // WebSocket. A `Harness` owns the runtime plus the receiver ends of its + // event channels so spawned senders never error. + // ---------------------------------------------------------------------- + + /// In-memory sink that records every `Message` written to it. All poll_* + /// hooks succeed immediately; `start_send` just pushes into a shared Vec. + #[derive(Clone)] + struct RecordingSink { + sent: Arc>>, + } + + impl RecordingSink { + fn new() -> Self { + Self { + sent: Arc::new(std::sync::Mutex::new(Vec::new())), + } + } + + /// All recorded messages decoded into `AgentMessage` (text frames only). + fn agent_messages(&self) -> Vec { + self.sent + .lock() + .unwrap() + .iter() + .filter_map(|m| match m { + Message::Text(t) => serde_json::from_str::(t.as_str()).ok(), + _ => None, + }) + .collect() + } + + fn sent_count(&self) -> usize { + self.sent.lock().unwrap().len() + } + } + + impl futures_util::Sink for RecordingSink { + type Error = tokio_tungstenite::tungstenite::Error; + + fn poll_ready( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Ready(Ok(())) + } + + fn start_send(self: std::pin::Pin<&mut Self>, item: Message) -> Result<(), Self::Error> { + self.sent.lock().unwrap().push(item); + Ok(()) + } + + fn poll_flush( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Ready(Ok(())) + } + + fn poll_close( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Ready(Ok(())) + } + } + + /// Owns the runtime plus the receiver halves of its event channels, + /// keeping them alive so background senders never see a closed channel. + struct Harness { + runtime: ConnectionRuntime, + events: ConnectionEvents, + } + + impl Harness { + /// Build a harness with the given capability bits. `file_cfg` lets + /// individual tests enable the file manager with a temp root. + fn new(caps: u32, file_cfg: FileConfig) -> Self { + let config = AgentConfig { + server_url: "http://127.0.0.1:9527".to_string(), + token: "t".to_string(), + enrollment_code: String::new(), + collector: CollectorConfig::default(), + log: LogConfig::default(), + file: file_cfg, + ip_change: IpChangeConfig::default(), + upgrade: UpgradeConfig::default(), + security: SecurityConfig::default(), + capabilities: CapabilitiesConfig::default(), + }; + let capabilities = CapabilityAuthority::fixed(caps); + let firewall_manager = Arc::new(FirewallManager::new(Arc::new(CliNftExecutor))); + let (runtime, events) = ConnectionRuntime::new(&config, capabilities, firewall_manager); + Self { runtime, events } + } + + /// Dispatch a single `ServerMessage` (as JSON) through the method. + async fn dispatch(&mut self, text: &str, sink: &mut RecordingSink) -> anyhow::Result<()> { + self.runtime.handle_server_message(text, sink).await + } + } + + /// All capability bits set — every success arm runs. + const ALL_CAPS: u32 = serverbee_common::constants::CAP_VALID_MASK; + + fn enabled_file_cfg(root: &std::path::Path) -> FileConfig { + FileConfig { + enabled: true, + root_paths: vec![root.to_string_lossy().to_string()], + ..FileConfig::default() + } + } + + #[tokio::test] + async fn test_blocklist_mutation_denied_without_firewall_capability() { + use serverbee_common::constants::CAP_FIREWALL_BLOCK; + // All caps except firewall block — simulates a revoked capability. + let caps = ALL_CAPS & !CAP_FIREWALL_BLOCK; + let mut h = Harness::new(caps, FileConfig::default()); + let mut sink = RecordingSink::new(); + + // A mutating variant must be dropped before reaching the firewall + // manager, so nothing is written back and no nft command runs. + h.dispatch(r#"{"type":"blocklist_remove","id":"x"}"#, &mut sink) + .await + .unwrap(); + assert_eq!( + sink.sent_count(), + 0, + "blocklist mutation must be ignored when CAP_FIREWALL_BLOCK is off" + ); + } + + #[tokio::test] + async fn test_dispatch_unparseable_text_is_ignored() { + let mut h = Harness::new(ALL_CAPS, FileConfig::default()); + let mut sink = RecordingSink::new(); + // Not valid JSON for ServerMessage — must be swallowed as Ok with no output. + h.dispatch("this is not json", &mut sink).await.unwrap(); + h.dispatch(r#"{"type":"nonexistent_variant"}"#, &mut sink) + .await + .unwrap(); + assert_eq!( + sink.sent_count(), + 0, + "unparseable input must not emit anything" + ); + } + + #[tokio::test] + async fn test_dispatch_ping_responds_with_pong() { + let mut h = Harness::new(ALL_CAPS, FileConfig::default()); + let mut sink = RecordingSink::new(); + h.dispatch(r#"{"type":"ping"}"#, &mut sink).await.unwrap(); + let msgs = sink.agent_messages(); + assert_eq!(msgs.len(), 1); + assert!(matches!(msgs[0], AgentMessage::Pong)); + } + + #[tokio::test] + async fn test_dispatch_ack_and_welcome_are_noops() { + let mut h = Harness::new(ALL_CAPS, FileConfig::default()); + let mut sink = RecordingSink::new(); + h.dispatch(r#"{"type":"ack","msg_id":"m1"}"#, &mut sink) + .await + .unwrap(); + h.dispatch( + r#"{"type":"welcome","server_id":"s","protocol_version":1,"report_interval":3}"#, + &mut sink, + ) + .await + .unwrap(); + assert_eq!( + sink.sent_count(), + 0, + "ack/welcome must not write to the sink" + ); + } + + #[tokio::test] + async fn test_dispatch_exec_denied_when_capability_absent() { + // CAP_EXEC missing -> a CapabilityDenied is pushed onto cmd_result_tx + // (NOT the sink). We drain the channel to assert. + let caps = ALL_CAPS & !serverbee_common::constants::CAP_EXEC; + let mut h = Harness::new(caps, FileConfig::default()); + let mut sink = RecordingSink::new(); + h.dispatch( + r#"{"type":"exec","task_id":"task-42","command":"true","timeout":1}"#, + &mut sink, + ) + .await + .unwrap(); + assert_eq!( + sink.sent_count(), + 0, + "denied exec writes to channel, not sink" + ); + let denied = h + .events + .cmd_result_rx + .recv() + .await + .expect("denied msg expected"); + match denied { + AgentMessage::CapabilityDenied { + msg_id, + capability, + reason, + .. + } => { + assert_eq!(msg_id, Some("task-42".to_string())); + assert_eq!(capability, "exec"); + assert_eq!(reason, CapabilityDeniedReason::AgentCapabilityDisabled); + } + other => panic!("expected CapabilityDenied, got {other:?}"), + } + } + + #[tokio::test] + async fn test_dispatch_exec_allowed_runs_and_emits_task_result() { + // `true` is a deterministic, always-present POSIX builtin/command. + let mut h = Harness::new(ALL_CAPS, FileConfig::default()); + let mut sink = RecordingSink::new(); + h.dispatch( + r#"{"type":"exec","task_id":"task-ok","command":"true","timeout":5}"#, + &mut sink, + ) + .await + .unwrap(); + // The execution is spawned; await its TaskResult on the channel. + let result = tokio::time::timeout(Duration::from_secs(10), h.events.cmd_result_rx.recv()) + .await + .expect("task did not complete in time") + .expect("TaskResult expected"); + match result { + AgentMessage::TaskResult { result, .. } => { + assert_eq!(result.task_id, "task-ok"); + assert_eq!(result.exit_code, 0); + } + other => panic!("expected TaskResult, got {other:?}"), + } + } + + #[tokio::test] + async fn test_dispatch_ping_tasks_sync_is_accepted() { + let mut h = Harness::new(ALL_CAPS, FileConfig::default()); + let mut sink = RecordingSink::new(); + // Empty task list — sync clears the manager; dispatcher returns Ok with + // no WS output. + h.dispatch(r#"{"type":"ping_tasks_sync","tasks":[]}"#, &mut sink) + .await + .unwrap(); + assert_eq!(sink.sent_count(), 0); + } + + #[tokio::test] + async fn test_dispatch_network_probe_sync_is_accepted() { + let mut h = Harness::new(ALL_CAPS, FileConfig::default()); + let mut sink = RecordingSink::new(); + h.dispatch( + r#"{"type":"network_probe_sync","targets":[],"interval":30,"packet_count":3}"#, + &mut sink, + ) + .await + .unwrap(); + assert_eq!(sink.sent_count(), 0); + } + + #[tokio::test] + async fn test_dispatch_terminal_lifecycle_without_capability() { + // CAP_TERMINAL off: open() routes to the denied event (no PTY spawned), + // and input/resize/close on a missing session are safe no-ops. The + // dispatcher must return Ok for each and never touch the sink. + let caps = ALL_CAPS & !serverbee_common::constants::CAP_TERMINAL; + let mut h = Harness::new(caps, FileConfig::default()); + let mut sink = RecordingSink::new(); + h.dispatch( + r#"{"type":"terminal_open","session_id":"s1","rows":24,"cols":80}"#, + &mut sink, + ) + .await + .unwrap(); + h.dispatch( + r#"{"type":"terminal_input","session_id":"s1","data":"aGk="}"#, + &mut sink, + ) + .await + .unwrap(); + h.dispatch( + r#"{"type":"terminal_resize","session_id":"s1","rows":30,"cols":100}"#, + &mut sink, + ) + .await + .unwrap(); + h.dispatch(r#"{"type":"terminal_close","session_id":"s1"}"#, &mut sink) + .await + .unwrap(); + assert_eq!( + sink.sent_count(), + 0, + "terminal control writes nothing to the WS sink" + ); + } + + #[tokio::test] + async fn test_dispatch_upgrade_denied_when_capability_absent() { + // CAP_UPGRADE off -> denied is written DIRECTLY to the sink (not the + // channel), unlike Exec. + let caps = ALL_CAPS & !serverbee_common::constants::CAP_UPGRADE; + let mut h = Harness::new(caps, FileConfig::default()); + let mut sink = RecordingSink::new(); + h.dispatch( + r#"{"type":"upgrade","version":"9.9.9","job_id":"j1"}"#, + &mut sink, + ) + .await + .unwrap(); + let msgs = sink.agent_messages(); + assert_eq!(msgs.len(), 1); + match &msgs[0] { + AgentMessage::CapabilityDenied { + capability, reason, .. + } => { + assert_eq!(capability, "upgrade"); + assert_eq!(*reason, CapabilityDeniedReason::AgentCapabilityDisabled); + } + other => panic!("expected CapabilityDenied, got {other:?}"), + } + } + + #[tokio::test] + async fn test_dispatch_traceroute_denied_when_capability_absent() { + let caps = ALL_CAPS & !serverbee_common::constants::CAP_PING_ICMP; + let mut h = Harness::new(caps, FileConfig::default()); + let mut sink = RecordingSink::new(); + h.dispatch( + r#"{"type":"traceroute","request_id":"r1","target":"example.com","max_hops":30}"#, + &mut sink, + ) + .await + .unwrap(); + assert_eq!( + sink.sent_count(), + 0, + "denied traceroute goes to the channel" + ); + let denied = h + .events + .cmd_result_rx + .recv() + .await + .expect("denied expected"); + match denied { + AgentMessage::CapabilityDenied { + msg_id, + capability, + reason, + .. + } => { + assert_eq!(msg_id, Some("r1".to_string())); + assert_eq!(capability, "ping_icmp"); + assert_eq!(reason, CapabilityDeniedReason::AgentCapabilityDisabled); + } + other => panic!("expected CapabilityDenied, got {other:?}"), + } + } + + #[tokio::test] + async fn test_dispatch_traceroute_invalid_target_is_rejected() { + // Capability present but target fails validation -> a completed + // TracerouteRoundUpdate with an error is emitted on the channel. No + // real traceroute subprocess is spawned. + let mut h = Harness::new(ALL_CAPS, FileConfig::default()); + let mut sink = RecordingSink::new(); + h.dispatch( + r#"{"type":"traceroute","request_id":"r2","target":"bad target with spaces; rm -rf","max_hops":30}"#, + &mut sink, + ) + .await + .unwrap(); + assert_eq!(sink.sent_count(), 0); + let msg = h + .events + .cmd_result_rx + .recv() + .await + .expect("update expected"); + match msg { + AgentMessage::TracerouteRoundUpdate { + request_id, + completed, + error, + .. + } => { + assert_eq!(request_id, "r2"); + assert!(completed); + assert!(error.is_some(), "invalid target must carry an error"); + } + other => panic!("expected TracerouteRoundUpdate, got {other:?}"), + } + } + + #[tokio::test] + async fn test_dispatch_file_ops_denied_when_capability_absent() { + // CAP_FILE off -> each file op replies with a disabled error frame on + // the sink (capability-absent branch). + let caps = ALL_CAPS & !serverbee_common::constants::CAP_FILE; + let mut h = Harness::new(caps, FileConfig::default()); + let mut sink = RecordingSink::new(); + + h.dispatch( + r#"{"type":"file_list","msg_id":"m1","path":"/tmp"}"#, + &mut sink, + ) + .await + .unwrap(); + h.dispatch( + r#"{"type":"file_stat","msg_id":"m2","path":"/tmp"}"#, + &mut sink, + ) + .await + .unwrap(); + h.dispatch( + r#"{"type":"file_read","msg_id":"m3","path":"/tmp/x","max_size":1024}"#, + &mut sink, + ) + .await + .unwrap(); + h.dispatch( + r#"{"type":"file_write","msg_id":"m4","path":"/tmp/x","content":"aGk="}"#, + &mut sink, + ) + .await + .unwrap(); + h.dispatch( + r#"{"type":"file_delete","msg_id":"m5","path":"/tmp/x","recursive":false}"#, + &mut sink, + ) + .await + .unwrap(); + h.dispatch( + r#"{"type":"file_mkdir","msg_id":"m6","path":"/tmp/d"}"#, + &mut sink, + ) + .await + .unwrap(); + h.dispatch( + r#"{"type":"file_move","msg_id":"m7","from":"/tmp/a","to":"/tmp/b"}"#, + &mut sink, + ) + .await + .unwrap(); + h.dispatch( + r#"{"type":"file_download_start","transfer_id":"t1","path":"/tmp/x"}"#, + &mut sink, + ) + .await + .unwrap(); + h.dispatch( + r#"{"type":"file_upload_start","transfer_id":"t2","path":"/tmp/x","size":4}"#, + &mut sink, + ) + .await + .unwrap(); + h.dispatch( + r#"{"type":"file_upload_chunk","transfer_id":"t3","offset":0,"data":"aGk="}"#, + &mut sink, + ) + .await + .unwrap(); + h.dispatch( + r#"{"type":"file_upload_end","transfer_id":"t4"}"#, + &mut sink, + ) + .await + .unwrap(); + + let msgs = sink.agent_messages(); + // 11 dispatches each produce exactly one response frame. + assert_eq!(msgs.len(), 11, "each denied file op emits one frame"); + // Spot-check representative variants carry the disabled error. + assert!(msgs.iter().any(|m| matches!( + m, + AgentMessage::FileListResult { error: Some(e), .. } if e.contains("disabled") + ))); + assert!(msgs.iter().any(|m| matches!( + m, + AgentMessage::FileOpResult { success: false, error: Some(e), .. } if e.contains("disabled") + ))); + assert!(msgs.iter().any(|m| matches!( + m, + AgentMessage::FileDownloadError { error, .. } if error.contains("disabled") + ))); + assert!(msgs.iter().any(|m| matches!( + m, + AgentMessage::FileUploadError { error, .. } if error.contains("disabled") + ))); + } + + #[tokio::test] + async fn test_dispatch_file_download_cancel_is_silent_noop() { + // FileDownloadCancel has no capability gate and no response; it just + // calls cancel_download. Cancelling an unknown transfer is a no-op. + let mut h = Harness::new(ALL_CAPS, FileConfig::default()); + let mut sink = RecordingSink::new(); + h.dispatch( + r#"{"type":"file_download_cancel","transfer_id":"nope"}"#, + &mut sink, + ) + .await + .unwrap(); + assert_eq!(sink.sent_count(), 0); + } + + #[tokio::test] + async fn test_dispatch_file_ops_success_with_enabled_manager() { + // File manager enabled with a real temp root: exercise the success + // branches (mkdir -> write -> list -> stat -> read -> move -> delete). + let tmp = tempfile::tempdir().unwrap(); + let root = std::fs::canonicalize(tmp.path()).unwrap(); + let cfg = enabled_file_cfg(&root); + let mut h = Harness::new(ALL_CAPS, cfg); + let mut sink = RecordingSink::new(); + + let sub = root.join("sub"); + let file_a = sub.join("a.txt"); + let file_b = sub.join("b.txt"); + let mkdir = format!( + r#"{{"type":"file_mkdir","msg_id":"mk","path":"{}"}}"#, + sub.to_string_lossy() + ); + h.dispatch(&mkdir, &mut sink).await.unwrap(); + + // `validate_path` canonicalizes, which requires the target to already + // exist — mirror the file_manager's own tests by pre-creating an empty + // file so the write overwrites it. + std::fs::write(&file_a, "").unwrap(); + + // base64("hi") == "aGk=" + let write = format!( + r#"{{"type":"file_write","msg_id":"w","path":"{}","content":"aGk="}}"#, + file_a.to_string_lossy() + ); + h.dispatch(&write, &mut sink).await.unwrap(); + + let list = format!( + r#"{{"type":"file_list","msg_id":"ls","path":"{}"}}"#, + sub.to_string_lossy() + ); + h.dispatch(&list, &mut sink).await.unwrap(); + + let stat = format!( + r#"{{"type":"file_stat","msg_id":"st","path":"{}"}}"#, + file_a.to_string_lossy() + ); + h.dispatch(&stat, &mut sink).await.unwrap(); + + let read = format!( + r#"{{"type":"file_read","msg_id":"rd","path":"{}","max_size":1024}}"#, + file_a.to_string_lossy() + ); + h.dispatch(&read, &mut sink).await.unwrap(); + + let mv = format!( + r#"{{"type":"file_move","msg_id":"mv","from":"{}","to":"{}"}}"#, + file_a.to_string_lossy(), + file_b.to_string_lossy() + ); + h.dispatch(&mv, &mut sink).await.unwrap(); + + let del = format!( + r#"{{"type":"file_delete","msg_id":"del","path":"{}","recursive":false}}"#, + file_b.to_string_lossy() + ); + h.dispatch(&del, &mut sink).await.unwrap(); + + let msgs = sink.agent_messages(); + assert_eq!(msgs.len(), 7, "seven file ops, seven responses"); + + // mkdir succeeded + assert!(msgs.iter().any(|m| matches!( + m, + AgentMessage::FileOpResult { msg_id, success: true, .. } if msg_id == "mk" + ))); + // write succeeded + assert!(msgs.iter().any(|m| matches!( + m, + AgentMessage::FileOpResult { msg_id, success: true, .. } if msg_id == "w" + ))); + // list returned at least the written file, no error + assert!(msgs.iter().any(|m| matches!( + m, + AgentMessage::FileListResult { msg_id, error: None, entries, .. } + if msg_id == "ls" && entries.iter().any(|e| e.name == "a.txt") + ))); + // stat found the entry + assert!(msgs.iter().any(|m| matches!( + m, + AgentMessage::FileStatResult { msg_id, entry: Some(_), error: None } if msg_id == "st" + ))); + // read returned the base64 content of "hi" + assert!(msgs.iter().any(|m| matches!( + m, + AgentMessage::FileReadResult { msg_id, content: Some(c), error: None } + if msg_id == "rd" && c == "aGk=" + ))); + // move + delete succeeded + assert!(msgs.iter().any(|m| matches!( + m, + AgentMessage::FileOpResult { msg_id, success: true, .. } if msg_id == "mv" + ))); + assert!(msgs.iter().any(|m| matches!( + m, + AgentMessage::FileOpResult { msg_id, success: true, .. } if msg_id == "del" + ))); + } + + #[tokio::test] + async fn test_dispatch_file_upload_success_round_trip() { + // Enabled manager: start -> chunk -> end upload, all on the sink. + let tmp = tempfile::tempdir().unwrap(); + let root = std::fs::canonicalize(tmp.path()).unwrap(); + let cfg = enabled_file_cfg(&root); + let mut h = Harness::new(ALL_CAPS, cfg); + let mut sink = RecordingSink::new(); + + let dest = root.join("up.bin"); + // "hi" -> base64 "aGk=" -> 2 bytes + let start = format!( + r#"{{"type":"file_upload_start","transfer_id":"u1","path":"{}","size":2}}"#, + dest.to_string_lossy() + ); + h.dispatch(&start, &mut sink).await.unwrap(); + h.dispatch( + r#"{"type":"file_upload_chunk","transfer_id":"u1","offset":0,"data":"aGk="}"#, + &mut sink, + ) + .await + .unwrap(); + h.dispatch( + r#"{"type":"file_upload_end","transfer_id":"u1"}"#, + &mut sink, + ) + .await + .unwrap(); + + let msgs = sink.agent_messages(); + assert_eq!(msgs.len(), 3); + // start ack at offset 0 + assert!(matches!( + &msgs[0], + AgentMessage::FileUploadAck { transfer_id, offset: 0 } if transfer_id == "u1" + )); + // chunk ack advances offset to 2 + assert!(matches!( + &msgs[1], + AgentMessage::FileUploadAck { transfer_id, offset: 2 } if transfer_id == "u1" + )); + // upload complete + assert!(matches!( + &msgs[2], + AgentMessage::FileUploadComplete { transfer_id } if transfer_id == "u1" + )); + // bytes actually landed on disk + assert_eq!(std::fs::read(&dest).unwrap(), b"hi"); + } + + #[tokio::test] + async fn test_dispatch_docker_start_stats_unavailable_emits_unavailable() { + // docker_manager is None -> DockerStartStats replies DockerUnavailable + // and leaves the stats interval unset. + let mut h = Harness::new(ALL_CAPS, FileConfig::default()); + let mut sink = RecordingSink::new(); + h.dispatch( + r#"{"type":"docker_start_stats","interval_secs":2}"#, + &mut sink, + ) + .await + .unwrap(); + let msgs = sink.agent_messages(); + assert_eq!(msgs.len(), 1); + assert!(matches!( + &msgs[0], + AgentMessage::DockerUnavailable { msg_id: None } + )); + assert!(h.runtime.docker_stats_interval.is_none()); + } + + #[tokio::test] + async fn test_dispatch_docker_stop_stats_clears_interval() { + let mut h = Harness::new(ALL_CAPS, FileConfig::default()); + // Pre-seed an interval so we can observe it being cleared. + h.runtime.docker_stats_interval = Some(tokio::time::interval(Duration::from_secs(60))); + let mut sink = RecordingSink::new(); + h.dispatch(r#"{"type":"docker_stop_stats"}"#, &mut sink) + .await + .unwrap(); + assert_eq!(sink.sent_count(), 0); + assert!(h.runtime.docker_stats_interval.is_none()); + } + + #[tokio::test] + async fn test_dispatch_docker_request_unavailable_carries_msg_id() { + // Request variants with docker_manager None reply DockerUnavailable + // echoing the request's msg_id. + let mut h = Harness::new(ALL_CAPS, FileConfig::default()); + let mut sink = RecordingSink::new(); + h.dispatch( + r#"{"type":"docker_list_containers","msg_id":"req-1"}"#, + &mut sink, + ) + .await + .unwrap(); + let msgs = sink.agent_messages(); + assert_eq!(msgs.len(), 1); + assert!(matches!( + &msgs[0], + AgentMessage::DockerUnavailable { msg_id: Some(id) } if id == "req-1" + )); + + // An event variant with no msg_id replies with msg_id: None. + let mut sink2 = RecordingSink::new(); + h.dispatch(r#"{"type":"docker_events_start"}"#, &mut sink2) + .await + .unwrap(); + let msgs2 = sink2.agent_messages(); + assert_eq!(msgs2.len(), 1); + assert!(matches!( + &msgs2[0], + AgentMessage::DockerUnavailable { msg_id: None } + )); + } + + #[tokio::test] + async fn test_dispatch_blocklist_reset_returns_ack() { + // FirewallManager uses the real CliNftExecutor. On a host without `nft` + // (macOS CI) BlocklistReset deterministically fails the wipe but still + // returns a BlocklistResetAck reply, which the dispatcher forwards. + let mut h = Harness::new(ALL_CAPS, FileConfig::default()); + let mut sink = RecordingSink::new(); + h.dispatch(r#"{"type":"blocklist_reset"}"#, &mut sink) + .await + .unwrap(); + let msgs = sink.agent_messages(); + assert_eq!(msgs.len(), 1, "reset always produces an ack"); + assert!(matches!(&msgs[0], AgentMessage::BlocklistResetAck { .. })); + } + + #[tokio::test] + async fn test_dispatch_ip_quality_sync_and_run_now_respect_capability() { + // With CAP_IP_QUALITY present, sync/run_now are accepted (no WS output). + let mut h = Harness::new(ALL_CAPS, FileConfig::default()); + let mut sink = RecordingSink::new(); + h.dispatch( + r#"{"type":"ip_quality_sync","services":[],"interval_hours":12}"#, + &mut sink, + ) + .await + .unwrap(); + h.dispatch(r#"{"type":"ip_quality_run_now"}"#, &mut sink) + .await + .unwrap(); + assert_eq!(sink.sent_count(), 0); + + // Without CAP_IP_QUALITY, both are silently ignored as well. + let caps = ALL_CAPS & !serverbee_common::constants::CAP_IP_QUALITY; + let mut h2 = Harness::new(caps, FileConfig::default()); + let mut sink2 = RecordingSink::new(); + h2.dispatch( + r#"{"type":"ip_quality_sync","services":[],"interval_hours":6}"#, + &mut sink2, + ) + .await + .unwrap(); + h2.dispatch(r#"{"type":"ip_quality_run_now"}"#, &mut sink2) + .await + .unwrap(); + assert_eq!(sink2.sent_count(), 0); + } + + #[tokio::test] + async fn test_dispatch_blocklist_sync_add_remove_forward_acks() { + // Sync/Add/Remove all route into FirewallManager and forward its + // BlocklistAck reply over the WS sink. On a host without `nft` the + // apply fails but the manager still returns an ack (with Failed state), + // so the dispatcher always emits exactly one frame per request. + let mut h = Harness::new(ALL_CAPS, FileConfig::default()); + + // Full-state sync with one entry. + let mut sink = RecordingSink::new(); + h.dispatch( + r#"{"type":"blocklist_sync","entries":[{"id":"e1","target":"1.2.3.4/32","family":4}]}"#, + &mut sink, + ) + .await + .unwrap(); + let msgs = sink.agent_messages(); + assert_eq!(msgs.len(), 1, "sync emits one ack frame"); + assert!(matches!(&msgs[0], AgentMessage::BlocklistAck { .. })); + + // Incremental add. + let mut sink = RecordingSink::new(); + h.dispatch( + r#"{"type":"blocklist_add","entry":{"id":"e2","target":"5.6.7.8/32","family":4}}"#, + &mut sink, + ) + .await + .unwrap(); + let msgs = sink.agent_messages(); + assert_eq!(msgs.len(), 1, "add emits one ack frame"); + assert!(matches!(&msgs[0], AgentMessage::BlocklistAck { .. })); + + // Incremental remove of an unknown id still produces a single-item ack. + let mut sink = RecordingSink::new(); + h.dispatch(r#"{"type":"blocklist_remove","id":"e2"}"#, &mut sink) + .await + .unwrap(); + let msgs = sink.agent_messages(); + assert_eq!(msgs.len(), 1, "remove emits one ack frame"); + assert!(matches!(&msgs[0], AgentMessage::BlocklistAck { .. })); + } + + #[tokio::test] + async fn test_dispatch_file_ops_disabled_manager_replies_disabled_even_with_capability() { + // CAP_FILE present but the manager is disabled (default FileConfig has + // enabled=false). The `!file_manager.is_enabled()` half of the guard + // must still short-circuit with the disabled error, independent of caps. + let mut h = Harness::new(ALL_CAPS, FileConfig::default()); + assert!( + !h.runtime.file_manager.is_enabled(), + "default file manager is disabled" + ); + let mut sink = RecordingSink::new(); + h.dispatch( + r#"{"type":"file_list","msg_id":"m1","path":"/tmp"}"#, + &mut sink, + ) + .await + .unwrap(); + h.dispatch( + r#"{"type":"file_write","msg_id":"m2","path":"/tmp/x","content":"aGk="}"#, + &mut sink, + ) + .await + .unwrap(); + let msgs = sink.agent_messages(); + assert_eq!(msgs.len(), 2, "each op replies once even with cap present"); + assert!(msgs.iter().any(|m| matches!( + m, + AgentMessage::FileListResult { error: Some(e), .. } if e.contains("disabled") + ))); + assert!(msgs.iter().any(|m| matches!( + m, + AgentMessage::FileOpResult { success: false, error: Some(e), .. } if e.contains("disabled") + ))); + } + + #[tokio::test] + async fn test_dispatch_upgrade_already_running_emits_failure_on_channel() { + // Force the global single-flight latch to "in progress", then dispatch + // an Upgrade with the capability present. The duplicate must be rejected + // with an UpgradeResult error on the cmd channel (not the WS sink), and + // the latch must be left untouched (still true) for the real holder. + UPGRADE_IN_PROGRESS.store(true, Ordering::SeqCst); + // Ensure we always release the global latch so other tests aren't poisoned. + struct Guard; + impl Drop for Guard { + fn drop(&mut self) { + UPGRADE_IN_PROGRESS.store(false, Ordering::SeqCst); + } + } + let _guard = Guard; + + let mut h = Harness::new(ALL_CAPS, FileConfig::default()); + let mut sink = RecordingSink::new(); + h.dispatch( + r#"{"type":"upgrade","version":"9.9.9","job_id":"dup-job"}"#, + &mut sink, + ) + .await + .unwrap(); + // Nothing is written to the WS sink for the duplicate case. + assert_eq!( + sink.sent_count(), + 0, + "duplicate upgrade writes to channel, not sink" + ); + let msg = tokio::time::timeout(Duration::from_secs(5), h.events.cmd_result_rx.recv()) + .await + .expect("failure msg expected in time") + .expect("UpgradeResult expected"); + match msg { + AgentMessage::UpgradeResult { + job_id, + target_version, + stage, + error, + .. + } => { + assert_eq!(job_id, Some("dup-job".to_string())); + assert_eq!(target_version, "9.9.9"); + assert_eq!(stage, UpgradeStage::Downloading); + assert!(error.contains("already running")); + } + other => panic!("expected UpgradeResult, got {other:?}"), + } + } +} From 6d023e6b6b442364a13af4e256f95fe88877a527 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 14:10:04 +0800 Subject: [PATCH 09/44] refactor(web): own server detail navigation policy in one module Tab catalogs, search validation, range catalogs and resolution, the legacy hours-to-key conversion, and the show_network/show_server_detail toggle combos were each restated at their call sites, so the admin detail route, the public detail route, the shared detail content, the network overview, and both public network routes all had to agree by copy. lib/server-detail-nav.ts now owns the whole policy: - SERVER_DETAIL_TABS / PUBLIC_SERVER_DETAIL_TABS plus the two validateSearch parsers move out of the route files - ADMIN/PUBLIC_TIME_RANGES, rangesForVariant, and resolveRange move out of server-detail-content - the legacy /network/$serverId redirect uses the shared legacyNetworkRangeToRangeKey conversion - publicNetworkHome(config) names the ADR-0001 decision (tab / standalone fallback / hidden) once; both public network routes and the overview deep-links consume the verdict instead of re-deriving it from raw toggles Behavior-preserving: the standalone public network page still renders only for show_server_detail=false + show_network=true deployments. --- .../status/server-detail-content.tsx | 30 +--- apps/web/src/lib/server-detail-nav.ts | 131 ++++++++++++++++++ .../src/routes/_authed/network/$serverId.tsx | 15 +- .../src/routes/_authed/servers/$id-page.tsx | 2 +- apps/web/src/routes/_authed/servers/$id.tsx | 9 +- .../_authed/servers/server-detail-search.ts | 2 - .../src/routes/status.network.$serverId.tsx | 15 +- apps/web/src/routes/status.network.index.tsx | 13 +- .../src/routes/status.server.$serverId.tsx | 27 ++-- 9 files changed, 161 insertions(+), 83 deletions(-) create mode 100644 apps/web/src/lib/server-detail-nav.ts delete mode 100644 apps/web/src/routes/_authed/servers/server-detail-search.ts diff --git a/apps/web/src/components/status/server-detail-content.tsx b/apps/web/src/components/status/server-detail-content.tsx index 86f297799..1f57f358f 100644 --- a/apps/web/src/components/status/server-detail-content.tsx +++ b/apps/web/src/components/status/server-detail-content.tsx @@ -26,16 +26,10 @@ import type { } from '@/lib/api-schema' import { buildMergedDiskIoSeries, buildPerDiskIoSeries } from '@/lib/disk-io' import { type ServerMetrics, useLiveServers } from '@/lib/server-catalog' +import { rangesForVariant, resolveRange, type TimeRange } from '@/lib/server-detail-nav' import { cn, formatBytes } from '@/lib/utils' import { computeAggregateUptime } from '@/lib/widget-helpers' -interface TimeRange { - hours: number - interval: string - key: string - label: string -} - interface GpuRecordAggregated { gpu_usage_avg: number mem_total_avg: number @@ -44,20 +38,6 @@ interface GpuRecordAggregated { time: string } -const ADMIN_TIME_RANGES: TimeRange[] = [ - { key: 'realtime', label: 'range_realtime', hours: 0, interval: 'realtime' }, - { key: '1h', label: 'range_1h', hours: 1, interval: 'raw' }, - { key: '6h', label: 'range_6h', hours: 6, interval: 'raw' }, - { key: '24h', label: 'range_24h', hours: 24, interval: 'raw' }, - { key: '7d', label: 'range_7d', hours: 168, interval: 'hourly' }, - { key: '30d', label: 'range_30d', hours: 720, interval: 'hourly' } -] - -// Public variant cannot rely on WS-driven realtime metrics, so realtime is -// dropped; everything else mirrors the admin range options because the -// public metrics endpoint accepts the same `interval` query parameter. -const PUBLIC_TIME_RANGES: TimeRange[] = ADMIN_TIME_RANGES.filter((r) => r.key !== 'realtime') - export interface ServerDetailContentProps { /** Currently selected detail tab. When provided (admin), the tabs become * URL-controlled; the public surface omits it and stays uncontrolled. */ @@ -85,12 +65,6 @@ function isAdminServer(server: ServerResponse | PublicServerDetail): server is S return 'ipv4' in server } -function resolveRange(rangeKey: string | undefined, ranges: TimeRange[]) { - const idx = ranges.findIndex((tr) => tr.key === rangeKey) - const rangeIndex = idx >= 0 ? idx : 0 - return { range: ranges[rangeIndex], rangeIndex } -} - function buildIsoWindow(hours: number) { const now = new Date() return { @@ -221,7 +195,7 @@ export function ServerDetailContent(props: ServerDetailContentProps) { const { t } = useTranslation('servers') const isPublic = variant === 'public' const isAdminVariant = !isPublic - const ranges = isPublic ? PUBLIC_TIME_RANGES : ADMIN_TIME_RANGES + const ranges = rangesForVariant(variant) const { range, rangeIndex } = resolveRange(rangeKey, ranges) const isRealtime = range.key === 'realtime' diff --git a/apps/web/src/lib/server-detail-nav.ts b/apps/web/src/lib/server-detail-nav.ts new file mode 100644 index 000000000..05746694c --- /dev/null +++ b/apps/web/src/lib/server-detail-nav.ts @@ -0,0 +1,131 @@ +/** + * Navigation policy for the server detail surface, admin and public. + * + * One module owns which tabs a detail URL may name, which historical ranges + * each variant offers, how a range key resolves, and where the public network + * detail experience lives for a given status-page config (ADR-0001: it is a + * server-detail tab whenever `show_server_detail` allows one, and the + * standalone `/status/network/$serverId` page only survives as a fallback for + * `show_server_detail=false` + `show_network=true` deployments). + */ + +// --- Tab catalogs ------------------------------------------------------------ + +export const SERVER_DETAIL_TABS = ['metrics', 'network', 'traffic', 'security', 'ip-quality'] as const +export type ServerDetailTab = (typeof SERVER_DETAIL_TABS)[number] + +export const PUBLIC_SERVER_DETAIL_TABS = ['metrics', 'network'] as const +export type PublicServerDetailTab = (typeof PUBLIC_SERVER_DETAIL_TABS)[number] + +/** Admin detail search: `range` always present (default realtime); `tab` kept + * optional so old `/servers/$id` links (range only) stay valid. */ +export function parseServerDetailSearch(search: Record): { range: string; tab?: ServerDetailTab } { + const range = (search.range as string) || 'realtime' + return SERVER_DETAIL_TABS.includes(search.tab as ServerDetailTab) + ? { range, tab: search.tab as ServerDetailTab } + : { range } +} + +/** Public detail search: both params optional; unknown tabs are dropped. */ +export function parsePublicServerDetailSearch(search: Record): { + range?: string + tab?: PublicServerDetailTab +} { + return { + range: typeof search.range === 'string' ? search.range : undefined, + ...(PUBLIC_SERVER_DETAIL_TABS.includes(search.tab as PublicServerDetailTab) + ? { tab: search.tab as PublicServerDetailTab } + : {}) + } +} + +/** A `?tab=network` deep link falls back to metrics when the network trigger + * is hidden, instead of selecting a tab that does not exist. */ +export function resolvePublicActiveTab( + tab: PublicServerDetailTab | undefined, + networkTabAvailable: boolean +): PublicServerDetailTab { + return networkTabAvailable && tab === 'network' ? 'network' : 'metrics' +} + +// --- Range catalogs ---------------------------------------------------------- + +export interface TimeRange { + hours: number + interval: string + key: string + label: string +} + +export const ADMIN_TIME_RANGES: TimeRange[] = [ + { key: 'realtime', label: 'range_realtime', hours: 0, interval: 'realtime' }, + { key: '1h', label: 'range_1h', hours: 1, interval: 'raw' }, + { key: '6h', label: 'range_6h', hours: 6, interval: 'raw' }, + { key: '24h', label: 'range_24h', hours: 24, interval: 'raw' }, + { key: '7d', label: 'range_7d', hours: 168, interval: 'hourly' }, + { key: '30d', label: 'range_30d', hours: 720, interval: 'hourly' } +] + +// Public variant cannot rely on WS-driven realtime metrics, so realtime is +// dropped; everything else mirrors the admin range options because the +// public metrics endpoint accepts the same `interval` query parameter. +export const PUBLIC_TIME_RANGES: TimeRange[] = ADMIN_TIME_RANGES.filter((r) => r.key !== 'realtime') + +export function rangesForVariant(variant: 'admin' | 'public'): TimeRange[] { + return variant === 'public' ? PUBLIC_TIME_RANGES : ADMIN_TIME_RANGES +} + +/** Resolve a range key against a catalog, falling back to the first entry so + * an unknown or absent key still renders a chart instead of nothing. */ +export function resolveRange(rangeKey: string | undefined, ranges: TimeRange[]) { + const idx = ranges.findIndex((tr) => tr.key === rangeKey) + const rangeIndex = idx >= 0 ? idx : 0 + return { range: ranges[rangeIndex], rangeIndex } +} + +/** The retired standalone network page counted hours in its `range` param + * ('1' | '6' | ...); detail URLs use metrics-style keys. Used by the legacy + * `/network/$serverId` redirect so old bookmarks land on the right window. */ +export function legacyNetworkRangeToRangeKey(range: string): string { + const mapping: Record = { + realtime: 'realtime', + '1': '1h', + '6': '6h', + '24': '24h', + '168': '7d', + '720': '30d' + } + return mapping[range] ?? 'realtime' +} + +// --- Public network home (ADR-0001) ------------------------------------------- + +interface PublicStatusToggles { + show_network?: boolean | null + show_server_detail?: boolean | null +} + +export type PublicNetworkHome = 'hidden' | 'standalone' | 'tab' + +/** + * Where the public network detail lives for this deployment. Returns + * `undefined` while the status config is still loading so callers can defer + * redirects instead of flashing the wrong page. + * + * - `tab` — inside `/status/server/$serverId?tab=network` (canonical home) + * - `standalone` — `/status/network/$serverId` fallback when the server + * detail page is disabled but network stays visible + * - `hidden` — network detail is not exposed at all + */ +export function publicNetworkHome(config: PublicStatusToggles | undefined): PublicNetworkHome | undefined { + if (config === undefined) { + return undefined + } + if (config.show_network === false) { + return 'hidden' + } + if (config.show_server_detail === false) { + return 'standalone' + } + return 'tab' +} diff --git a/apps/web/src/routes/_authed/network/$serverId.tsx b/apps/web/src/routes/_authed/network/$serverId.tsx index 325f5059f..5e66ba2fd 100644 --- a/apps/web/src/routes/_authed/network/$serverId.tsx +++ b/apps/web/src/routes/_authed/network/$serverId.tsx @@ -1,18 +1,9 @@ import { createFileRoute, redirect } from '@tanstack/react-router' +import { legacyNetworkRangeToRangeKey } from '@/lib/server-detail-nav' // The standalone network detail page merged into the server detail Network // tab. This route survives purely as a redirect so old bookmarks and links -// keep working. The old `range` param counted hours ('1' | '6' | ... ); -// the server detail page uses metrics-style keys. -const HOURS_TO_RANGE_KEY: Record = { - realtime: 'realtime', - '1': '1h', - '6': '6h', - '24': '24h', - '168': '7d', - '720': '30d' -} - +// keep working. export const Route = createFileRoute('/_authed/network/$serverId')({ validateSearch: (search: Record) => ({ range: (search.range as string) || 'realtime' @@ -21,7 +12,7 @@ export const Route = createFileRoute('/_authed/network/$serverId')({ throw redirect({ to: '/servers/$id', params: { id: params.serverId }, - search: { tab: 'network', range: HOURS_TO_RANGE_KEY[search.range] ?? 'realtime' }, + search: { tab: 'network', range: legacyNetworkRangeToRangeKey(search.range) }, replace: true }) } diff --git a/apps/web/src/routes/_authed/servers/$id-page.tsx b/apps/web/src/routes/_authed/servers/$id-page.tsx index 00d26545d..97041bb2c 100644 --- a/apps/web/src/routes/_authed/servers/$id-page.tsx +++ b/apps/web/src/routes/_authed/servers/$id-page.tsx @@ -20,9 +20,9 @@ import { api } from '@/lib/api-client' import type { ServerResponse } from '@/lib/api-schema' import { CAP_DOCKER, CAP_FILE, CAP_TERMINAL, getEffectiveCapabilityEnabled } from '@/lib/capabilities' import { useLiveServers } from '@/lib/server-catalog' +import type { ServerDetailTab } from '@/lib/server-detail-nav' import { formatBytes } from '@/lib/utils' import { useUpgradeJobsStore } from '@/stores/upgrade-jobs-store' -import type { ServerDetailTab } from './server-detail-search' const routeApi = getRouteApi('/_authed/servers/$id') diff --git a/apps/web/src/routes/_authed/servers/$id.tsx b/apps/web/src/routes/_authed/servers/$id.tsx index 3db47be94..3ee3aafe8 100644 --- a/apps/web/src/routes/_authed/servers/$id.tsx +++ b/apps/web/src/routes/_authed/servers/$id.tsx @@ -1,7 +1,7 @@ import { createFileRoute } from '@tanstack/react-router' import { lazy, Suspense } from 'react' import { Skeleton } from '@/components/ui/skeleton' -import { SERVER_DETAIL_TABS, type ServerDetailTab } from './server-detail-search' +import { parseServerDetailSearch } from '@/lib/server-detail-nav' const LazyServerDetailPage = lazy(() => import('./$id-page').then((m) => ({ default: m.ServerDetailPage }))) @@ -30,10 +30,5 @@ export const Route = createFileRoute('/_authed/servers/$id')({ // the same tab (e.g. Security) instead of resetting to Metrics. `tab` is optional // so existing links to /servers/$id (range only) stay valid; the page defaults to // the metrics tab when it is absent. - validateSearch: (search: Record): { range: string; tab?: ServerDetailTab } => { - const range = (search.range as string) || 'realtime' - return SERVER_DETAIL_TABS.includes(search.tab as ServerDetailTab) - ? { range, tab: search.tab as ServerDetailTab } - : { range } - } + validateSearch: parseServerDetailSearch }) diff --git a/apps/web/src/routes/_authed/servers/server-detail-search.ts b/apps/web/src/routes/_authed/servers/server-detail-search.ts deleted file mode 100644 index ab2bcb7f6..000000000 --- a/apps/web/src/routes/_authed/servers/server-detail-search.ts +++ /dev/null @@ -1,2 +0,0 @@ -export const SERVER_DETAIL_TABS = ['metrics', 'network', 'traffic', 'security', 'ip-quality'] as const -export type ServerDetailTab = (typeof SERVER_DETAIL_TABS)[number] diff --git a/apps/web/src/routes/status.network.$serverId.tsx b/apps/web/src/routes/status.network.$serverId.tsx index 2bf96672e..e2fc1793c 100644 --- a/apps/web/src/routes/status.network.$serverId.tsx +++ b/apps/web/src/routes/status.network.$serverId.tsx @@ -10,6 +10,7 @@ import { usePublicStatusConfig } from '@/hooks/use-public-status' import { api } from '@/lib/api-client' import type { PublicNetworkServerDetail } from '@/lib/api-schema' import { formatDateTime } from '@/lib/format' +import { publicNetworkHome } from '@/lib/server-detail-nav' export const Route = createFileRoute('/status/network/$serverId')({ component: PublicNetworkDetailPage @@ -33,9 +34,8 @@ function PublicNetworkDetailPage() { const navigate = useNavigate() const [anomalyOpen, setAnomalyOpen] = useState(false) - const networkEnabled = config?.show_network !== false - const detailEnabled = config?.show_server_detail !== false - const isStandalone = config !== undefined && networkEnabled && !detailEnabled + const networkHome = publicNetworkHome(config) + const isStandalone = networkHome === 'standalone' const { data, isLoading, error } = useQuery({ queryKey: ['public-status', 'network', serverId], @@ -46,14 +46,11 @@ function PublicNetworkDetailPage() { }) useEffect(() => { - if (!config) { - return - } - if (config.show_network === false) { + if (networkHome === 'hidden') { navigate({ to: '/status', replace: true }) return } - if (detailEnabled) { + if (networkHome === 'tab') { navigate({ to: '/status/server/$serverId', params: { serverId }, @@ -61,7 +58,7 @@ function PublicNetworkDetailPage() { replace: true }) } - }, [config, detailEnabled, navigate, serverId]) + }, [networkHome, navigate, serverId]) // Wait for the config before choosing between redirect and standalone so we // never flash the standalone page on sites that will redirect. diff --git a/apps/web/src/routes/status.network.index.tsx b/apps/web/src/routes/status.network.index.tsx index c1adc2223..d64a5451b 100644 --- a/apps/web/src/routes/status.network.index.tsx +++ b/apps/web/src/routes/status.network.index.tsx @@ -6,6 +6,7 @@ import { NetworkOverviewContent } from '@/components/status/network-overview-con import { usePublicStatusConfig } from '@/hooks/use-public-status' import { api } from '@/lib/api-client' import type { PublicNetworkOverview } from '@/lib/api-schema' +import { publicNetworkHome } from '@/lib/server-detail-nav' export const Route = createFileRoute('/status/network/')({ component: PublicNetworkOverviewPage @@ -17,22 +18,22 @@ function PublicNetworkOverviewPage() { const navigate = useNavigate() const [search, setSearch] = useState('') - const networkEnabled = config?.show_network !== false + const networkHome = publicNetworkHome(config) const { data, isLoading, error } = useQuery({ queryKey: ['public-status', 'network', 'overview'], queryFn: () => api.get('/api/status/network'), refetchInterval: 30_000, - enabled: networkEnabled, + enabled: networkHome !== 'hidden', retry: false }) useEffect(() => { - if (config && config.show_network === false) { + if (networkHome === 'hidden') { navigate({ to: '/status', replace: true }) } - }, [config, navigate]) + }, [networkHome, navigate]) - if (config?.show_network === false) { + if (networkHome === 'hidden') { return null } @@ -49,7 +50,7 @@ function PublicNetworkOverviewPage() { data={data?.servers ?? []} isLoading={isLoading} onSearchChange={setSearch} - publicDetailTabEnabled={config?.show_server_detail !== false} + publicDetailTabEnabled={networkHome !== 'standalone'} search={search} variant="public" /> diff --git a/apps/web/src/routes/status.server.$serverId.tsx b/apps/web/src/routes/status.server.$serverId.tsx index 0aebdf5dc..e8f95b53c 100644 --- a/apps/web/src/routes/status.server.$serverId.tsx +++ b/apps/web/src/routes/status.server.$serverId.tsx @@ -11,24 +11,17 @@ import { Skeleton } from '@/components/ui/skeleton' import { usePublicStatusConfig } from '@/hooks/use-public-status' import { api } from '@/lib/api-client' import type { PublicServerDetail } from '@/lib/api-schema' +import { + type PublicServerDetailTab, + parsePublicServerDetailSearch, + publicNetworkHome, + resolvePublicActiveTab +} from '@/lib/server-detail-nav' import { formatBytes } from '@/lib/utils' -const PUBLIC_SERVER_DETAIL_TABS = ['metrics', 'network'] as const -type PublicServerDetailTab = (typeof PUBLIC_SERVER_DETAIL_TABS)[number] - -interface PublicServerDetailSearch { - range?: string - tab?: PublicServerDetailTab -} - export const Route = createFileRoute('/status/server/$serverId')({ component: PublicServerDetailPage, - validateSearch: (search: Record): PublicServerDetailSearch => ({ - range: typeof search.range === 'string' ? search.range : undefined, - ...(PUBLIC_SERVER_DETAIL_TABS.includes(search.tab as PublicServerDetailTab) - ? { tab: search.tab as PublicServerDetailTab } - : {}) - }) + validateSearch: parsePublicServerDetailSearch }) function PublicServerDetailPage() { @@ -43,7 +36,7 @@ function PublicServerDetailPage() { const range = rangeParam const detailEnabled = config?.show_server_detail !== false - const networkEnabled = config?.show_network !== false + const networkEnabled = publicNetworkHome(config) !== 'hidden' const { data, isLoading, error } = useQuery({ queryKey: ['public-status', 'server', serverId], queryFn: () => api.get(`/api/status/servers/${serverId}`), @@ -83,9 +76,7 @@ function PublicServerDetailPage() { ) } - // If the network toggle is off, a `?tab=network` deep link falls back to - // metrics instead of selecting a trigger that does not exist. - const activeTab = networkEnabled && tabParam === 'network' ? 'network' : 'metrics' + const activeTab = resolvePublicActiveTab(tabParam, networkEnabled) return (
From 516fb918daf4ca2e640771f5cf4ad2b3e2095827 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 14:19:57 +0800 Subject: [PATCH 10/44] refactor(server): dedupe reconcile warn wrappers into the reconciler MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-up to the desired-state reconciler: ping/network/IP-quality routes each carried an identical warn-and-swallow wrapper, server.rs inlined the same shape, and the capability-change handler re-implemented reconcile_connection's four-domain loop by hand — a drift risk when a fifth domain lands. The reconciler now owns the demoted-to-warning variants (reconcile_connected_or_warn / reconcile_agent_or_warn) and the capability-change path calls reconcile_connection like the connect path does. --- crates/server/src/router/api/ip_quality.rs | 29 ++++++++++--------- crates/server/src/router/api/network_probe.rs | 24 +++++++-------- crates/server/src/router/api/ping.rs | 24 +++++++-------- crates/server/src/router/api/server.rs | 9 ++---- crates/server/src/router/ws/agent/security.rs | 29 +++++++------------ crates/server/src/service/agent_reconcile.rs | 26 +++++++++++++++++ 6 files changed, 79 insertions(+), 62 deletions(-) diff --git a/crates/server/src/router/api/ip_quality.rs b/crates/server/src/router/api/ip_quality.rs index ef10a4ab6..2a8b9c1cd 100644 --- a/crates/server/src/router/api/ip_quality.rs +++ b/crates/server/src/router/api/ip_quality.rs @@ -168,7 +168,10 @@ async fn create_service( Json(input): Json, ) -> Result>, AppError> { let service = IpQualityService::create_custom_service(&state.db, input).await?; - reconcile_ip_quality(&state).await; + state + .agent_desired_state + .reconcile_connected_or_warn(AgentDesiredStateDomain::IpQuality) + .await; ok(service) } @@ -190,7 +193,10 @@ async fn update_service( Json(input): Json, ) -> Result>, AppError> { let service = IpQualityService::update_service(&state.db, &id, input).await?; - reconcile_ip_quality(&state).await; + state + .agent_desired_state + .reconcile_connected_or_warn(AgentDesiredStateDomain::IpQuality) + .await; ok(service) } @@ -211,7 +217,10 @@ async fn delete_service( Path(id): Path, ) -> Result>, AppError> { IpQualityService::delete_service(&state.db, &id).await?; - reconcile_ip_quality(&state).await; + state + .agent_desired_state + .reconcile_connected_or_warn(AgentDesiredStateDomain::IpQuality) + .await; ok("ok") } @@ -232,19 +241,13 @@ async fn update_settings( Json(input): Json, ) -> Result>, AppError> { let setting = IpQualityService::update_setting(&state.db, input.check_interval_hours).await?; - reconcile_ip_quality(&state).await; + state + .agent_desired_state + .reconcile_connected_or_warn(AgentDesiredStateDomain::IpQuality) + .await; ok(setting) } -async fn reconcile_ip_quality(state: &AppState) { - if let Err(error) = state - .agent_desired_state - .reconcile_connected(AgentDesiredStateDomain::IpQuality) - .await - { - tracing::warn!(%error, "failed to reconcile IP quality desired state"); - } -} #[utoipa::path( post, diff --git a/crates/server/src/router/api/network_probe.rs b/crates/server/src/router/api/network_probe.rs index 79f1521f2..c1493713c 100644 --- a/crates/server/src/router/api/network_probe.rs +++ b/crates/server/src/router/api/network_probe.rs @@ -128,7 +128,10 @@ async fn update_target( Json(input): Json, ) -> Result>, AppError> { let target = NetworkProbeService::update_target(&state.db, &id, input).await?; - reconcile_network_probes(&state).await; + state + .agent_desired_state + .reconcile_connected_or_warn(AgentDesiredStateDomain::NetworkProbes) + .await; ok(target) } @@ -149,7 +152,10 @@ async fn delete_target( ) -> Result>, AppError> { // Delete the target (cascades records + configs + setting cleanup) NetworkProbeService::delete_target(&state.db, &id).await?; - reconcile_network_probes(&state).await; + state + .agent_desired_state + .reconcile_connected_or_warn(AgentDesiredStateDomain::NetworkProbes) + .await; ok("ok") } @@ -169,20 +175,14 @@ async fn update_setting( Json(input): Json, ) -> Result>, AppError> { NetworkProbeService::update_setting(&state.db, &input).await?; - reconcile_network_probes(&state).await; + state + .agent_desired_state + .reconcile_connected_or_warn(AgentDesiredStateDomain::NetworkProbes) + .await; ok(input) } -async fn reconcile_network_probes(state: &AppState) { - if let Err(error) = state - .agent_desired_state - .reconcile_connected(AgentDesiredStateDomain::NetworkProbes) - .await - { - tracing::warn!(%error, "failed to reconcile network probe desired state"); - } -} // --------------------------------------------------------------------------- // Per-server read handlers (mounted in server.rs) diff --git a/crates/server/src/router/api/ping.rs b/crates/server/src/router/api/ping.rs index b4a76bd7a..9e2ed4c65 100644 --- a/crates/server/src/router/api/ping.rs +++ b/crates/server/src/router/api/ping.rs @@ -82,7 +82,10 @@ async fn create_task( Json(input): Json, ) -> Result>, AppError> { let task = PingService::create(&state.db, input).await?; - reconcile_ping_tasks(&state).await; + state + .agent_desired_state + .reconcile_connected_or_warn(AgentDesiredStateDomain::PingTasks) + .await; ok(task) } @@ -105,7 +108,10 @@ async fn update_task( Json(input): Json, ) -> Result>, AppError> { let task = PingService::update(&state.db, &id, input).await?; - reconcile_ping_tasks(&state).await; + state + .agent_desired_state + .reconcile_connected_or_warn(AgentDesiredStateDomain::PingTasks) + .await; ok(task) } @@ -126,19 +132,13 @@ async fn delete_task( Path(id): Path, ) -> Result>, AppError> { PingService::delete(&state.db, &id).await?; - reconcile_ping_tasks(&state).await; + state + .agent_desired_state + .reconcile_connected_or_warn(AgentDesiredStateDomain::PingTasks) + .await; ok("ok") } -async fn reconcile_ping_tasks(state: &AppState) { - if let Err(error) = state - .agent_desired_state - .reconcile_connected(AgentDesiredStateDomain::PingTasks) - .await - { - tracing::warn!(%error, "failed to reconcile ping task desired state"); - } -} #[derive(Deserialize, utoipa::IntoParams)] pub struct RecordsQuery { diff --git a/crates/server/src/router/api/server.rs b/crates/server/src/router/api/server.rs index b0837d21e..a203732aa 100644 --- a/crates/server/src/router/api/server.rs +++ b/crates/server/src/router/api/server.rs @@ -1199,13 +1199,10 @@ async fn set_server_network_targets( ) -> Result>, AppError> { NetworkProbeService::set_server_targets(&state.db, &id, body.target_ids).await?; - if let Err(error) = state + state .agent_desired_state - .reconcile_agent(&id, AgentDesiredStateDomain::NetworkProbes) - .await - { - tracing::warn!(server_id = id, %error, "failed to reconcile network probe desired state"); - } + .reconcile_agent_or_warn(&id, AgentDesiredStateDomain::NetworkProbes) + .await; ok("ok") } diff --git a/crates/server/src/router/ws/agent/security.rs b/crates/server/src/router/ws/agent/security.rs index dcedc69cd..3b0dbe73c 100644 --- a/crates/server/src/router/ws/agent/security.rs +++ b/crates/server/src/router/ws/agent/security.rs @@ -8,7 +8,6 @@ use std::sync::Arc; use std::time::Duration; -use crate::service::agent_reconcile::AgentDesiredStateDomain; use crate::service::alert::AlertService; use crate::service::audit::AuditService; use crate::service::ip_quality::IpQualityService; @@ -187,24 +186,16 @@ pub(super) async fn on_capabilities_changed( tracing::error!("Failed to mirror capabilities for {server_id}: {e}"); } - for domain in [ - AgentDesiredStateDomain::PingTasks, - AgentDesiredStateDomain::NetworkProbes, - AgentDesiredStateDomain::IpQuality, - AgentDesiredStateDomain::Firewall, - ] { - if let Err(error) = state - .agent_desired_state - .reconcile_agent(server_id, domain) - .await - { - tracing::warn!( - server_id, - ?domain, - error = %error, - "capability-change desired-state reconcile failed" - ); - } + if let Err(error) = state + .agent_desired_state + .reconcile_connection(server_id) + .await + { + tracing::warn!( + server_id, + error = %error, + "capability-change desired-state reconcile was incomplete" + ); } // Resolve display name + originating IP for the audit trail. Neither diff --git a/crates/server/src/service/agent_reconcile.rs b/crates/server/src/service/agent_reconcile.rs index 1798bc8be..cd6b23ae1 100644 --- a/crates/server/src/service/agent_reconcile.rs +++ b/crates/server/src/service/agent_reconcile.rs @@ -159,6 +159,32 @@ impl AgentDesiredStateReconciler { } } + /// [`Self::reconcile_connected`], demoted to a warning. Mutation handlers + /// must not fail their HTTP request over a push failure — the agent's + /// state converges on its next connection reconcile. + pub async fn reconcile_connected_or_warn(&self, domain: AgentDesiredStateDomain) { + if let Err(error) = self.reconcile_connected(domain).await { + tracing::warn!( + ?domain, + error = %error, + "failed to reconcile connected agents' desired state" + ); + } + } + + /// [`Self::reconcile_agent`], demoted to a warning under the same + /// contract as [`Self::reconcile_connected_or_warn`]. + pub async fn reconcile_agent_or_warn(&self, server_id: &str, domain: AgentDesiredStateDomain) { + if let Err(error) = self.reconcile_agent(server_id, domain).await { + tracing::warn!( + server_id, + ?domain, + error = %error, + "failed to reconcile agent desired state" + ); + } + } + async fn reconcile_connected_ping_tasks(&self) -> Result<(), AppError> { let server_ids = self.agent_manager.connected_server_ids(); let tasks = self.load_ping_tasks().await?; From 38e050835d52def6c43137876079473f98e6bde5 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 14:21:42 +0800 Subject: [PATCH 11/44] docs: record desired-state reconcile fixes in the changelog Add Unreleased entries for the three user-visible behavior changes that rode in with the desired-state reconciler: capability changes resync everything the agent executes, network probe target edits reach connected agents, and a failed firewall blocklist reset keeps its pending state instead of being recorded as done. --- CHANGELOG.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4672dfe94..f39a6de93 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **Temporary capability grants now bound running work** -- A temporary `terminal` grant's expiry blocked new sessions but left live PTY sessions running indefinitely, and the security-events pipeline was started once at boot from the permanent capability set, so a temporary `security_events` grant could never start it (nor could a persisted grant enable it across a restart) and expiry never stopped it. Capability state is now owned by one process-wide authority on the agent: expiry or revocation closes live terminal sessions immediately, and the security pipeline starts and stops as the capability becomes effective or lapses +- **Capability changes now resync everything the agent executes for the server** -- When an agent reported a capability change mid-connection (a grant expiring, a config edit followed by `SIGHUP`), the server updated its mirror but never re-derived what the agent should be running: a host that revoked its ping capabilities kept executing the last-synced ping task list, and a revoked firewall capability left ServerBee's nftables entries in place until the next reconnect. Every capability change now triggers a full desired-state reconcile -- ping tasks, network probes, IP-quality services, and the firewall blocklist are re-filtered against the new capability set and pushed (or reset) immediately + +- **Editing a network probe target now reaches connected agents** -- Creating or deleting a global probe target and changing the probe settings pushed a fresh `NetworkProbeSync`, but editing an existing target's address or type did not, so connected agents kept probing the old target until they reconnected. Target updates now push the refreshed target list like every other probe mutation + +- **A failed firewall blocklist reset is no longer recorded as done** -- When an agent acknowledged a blocklist reset with `ok=false` (for example `nft` missing or permission denied), the server cleared its record of what the agent had applied anyway, so operators saw the host as wiped while the rules stayed live. A successful ack is now the only evidence that clears the per-server apply state; a failed ack keeps the last-known state visible and the wipe is retried on the next connection reconcile + ## [1.0.0-alpha.11] - 2026-07-03 ### Security From 7a039aca7e708961b9650a155e60d54548f10947 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 15:08:12 +0800 Subject: [PATCH 12/44] test(server): tolerate ip_quality_sync in first-connect noise filters The desired-state reconciler now pushes IpQualitySync on every fresh SystemInfo, but four mock-agent suites still panicked on any frame outside their hand-listed first-connect noise set, failing whenever the sync raced ahead of the docker/file command under test. Add ip_quality_sync to the noise filters, matching the suites that already tolerated it. --- crates/server/tests/agent_docker_extra.rs | 1 + crates/server/tests/agent_file_ops.rs | 1 + crates/server/tests/docker_integration.rs | 2 ++ crates/server/tests/router_file_extra.rs | 1 + 4 files changed, 5 insertions(+) diff --git a/crates/server/tests/agent_docker_extra.rs b/crates/server/tests/agent_docker_extra.rs index 0d1fb6ce8..7563713f9 100644 --- a/crates/server/tests/agent_docker_extra.rs +++ b/crates/server/tests/agent_docker_extra.rs @@ -76,6 +76,7 @@ fn is_first_connect_noise(msg_type: Option<&str>) -> bool { msg_type, Some("ping_tasks_sync") | Some("network_probe_sync") + | Some("ip_quality_sync") | Some("blocklist_reset") | Some("blocklist_sync") | Some("blocklist_add") diff --git a/crates/server/tests/agent_file_ops.rs b/crates/server/tests/agent_file_ops.rs index 5e12ae97e..f48b05a57 100644 --- a/crates/server/tests/agent_file_ops.rs +++ b/crates/server/tests/agent_file_ops.rs @@ -47,6 +47,7 @@ fn is_ignorable_push(msg_type: Option<&str>) -> bool { msg_type, Some("ping_tasks_sync") | Some("network_probe_sync") + | Some("ip_quality_sync") | Some("blocklist_reset") | Some("blocklist_sync") | Some("blocklist_add") diff --git a/crates/server/tests/docker_integration.rs b/crates/server/tests/docker_integration.rs index bc3daea83..583a59334 100644 --- a/crates/server/tests/docker_integration.rs +++ b/crates/server/tests/docker_integration.rs @@ -411,6 +411,7 @@ async fn test_docker_info_endpoint_requests_agent_when_cache_empty() { // docker flow under test and are ignored. Some("ping_tasks_sync") | Some("network_probe_sync") + | Some("ip_quality_sync") | Some("blocklist_reset") | Some("blocklist_sync") | Some("blocklist_add") @@ -788,6 +789,7 @@ async fn test_docker_unavailable_fails_pending_request_and_clears_feature_state( // docker flow under test and are ignored. Some("ping_tasks_sync") | Some("network_probe_sync") + | Some("ip_quality_sync") | Some("blocklist_reset") | Some("blocklist_sync") | Some("blocklist_add") diff --git a/crates/server/tests/router_file_extra.rs b/crates/server/tests/router_file_extra.rs index 8b51a38e5..edb52469b 100644 --- a/crates/server/tests/router_file_extra.rs +++ b/crates/server/tests/router_file_extra.rs @@ -54,6 +54,7 @@ fn is_ignorable_push(msg_type: Option<&str>) -> bool { msg_type, Some("ping_tasks_sync") | Some("network_probe_sync") + | Some("ip_quality_sync") | Some("blocklist_reset") | Some("blocklist_sync") | Some("blocklist_add") From 594b8f8808be6d925a5c201f33ec2d11e9731061 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 15:18:13 +0800 Subject: [PATCH 13/44] docs: document capability-change desired-state resync MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extend the 'What the server sees' list on the capabilities page (both languages): a capability change now re-syncs everything the agent executes for the server — ping tasks re-filtered, probe targets and IP-quality services re-pushed, firewall blocklist reset and only re-sent while firewall_block stays effective. --- apps/docs/content/docs/en/capabilities.mdx | 1 + apps/docs/content/docs/zh/capabilities.mdx | 1 + 2 files changed, 2 insertions(+) diff --git a/apps/docs/content/docs/en/capabilities.mdx b/apps/docs/content/docs/en/capabilities.mdx index 4af5d9201..9d6c1a023 100644 --- a/apps/docs/content/docs/en/capabilities.mdx +++ b/apps/docs/content/docs/en/capabilities.mdx @@ -147,6 +147,7 @@ When a grant becomes active (or expires/is revoked), the agent re-reports its ef - The change is mirrored into `servers.capabilities` and broadcast to browsers, so the UI updates in real time. - The transition is written to the [audit log](/en/docs/admin) as `capability_temporarily_granted`, `capability_grant_expired`, or `capability_grant_revoked`. - A temporary grant of a **high-risk** capability (`terminal`, `exec`, `file`, or `docker`) also evaluates the event-driven `capability_grant_detected` [alert rule](/en/docs/alerts). Expiry, revoke, and low-risk grants are audited but do not fire an alert. +- Everything the agent executes for the server is re-synced against the new capability set: ping tasks are re-filtered, network probe targets and IP-quality services are re-pushed, and the firewall blocklist is reset (and re-sent only while `firewall_block` is still effective). Revoking a ping capability stops those probes immediately, and revoking `firewall_block` removes ServerBee's nftables rules from the host without waiting for a reconnect. A grant issued while the agent is **disconnected** from the server is applied locally and shows up on reconnect, but it does **not** produce a `granted` audit entry or alert. The server only sees the post-reconnect capability state, not the moment of transition, so it cannot distinguish a freshly-granted capability from one that was already active. Plan around this if you rely on the alert as a tripwire. diff --git a/apps/docs/content/docs/zh/capabilities.mdx b/apps/docs/content/docs/zh/capabilities.mdx index a3d8aaddd..3090af676 100644 --- a/apps/docs/content/docs/zh/capabilities.mdx +++ b/apps/docs/content/docs/zh/capabilities.mdx @@ -147,6 +147,7 @@ serverbee-agent grants - 变更会镜像到 `servers.capabilities` 并广播给浏览器,UI 实时更新。 - 该转换会写入 [审计日志](/zh/docs/admin),动作为 `capability_temporarily_granted`、`capability_grant_expired` 或 `capability_grant_revoked`。 - 对**高危**能力(`terminal`、`exec`、`file`、`docker`)的临时授予还会评估事件驱动的 `capability_grant_detected` [告警规则](/zh/docs/alerts)。过期、撤销以及低风险授予会被审计,但不触发告警。 +- Agent 为该服务器执行的全部工作会按新的能力集重新同步:ping 任务重新过滤、网络探测目标与 IP 质量服务重新下发、防火墙黑名单先重置(仅当 `firewall_block` 仍生效时才会重新推送)。撤销 ping 能力会立即停止对应探测;撤销 `firewall_block` 会立刻移除 ServerBee 在主机上的 nftables 规则,无需等待重连。 在 Agent 与 Server **断连**期间发起的授予,会在本地生效并在重连后展示出来,但**不会**产生 `granted` 审计记录或告警。Server 只能看到重连后的能力状态,看不到转换发生的瞬间,因此无法区分「刚刚授予」与「本就生效」的能力。若你把该告警当作触发器(tripwire)使用,请考虑这一点。 From a985c689fa8bd84318be0271af1e462d1d109009 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 15:53:24 +0800 Subject: [PATCH 14/44] refactor(server): dedupe desired-state reconciler internals Map each domain to its provider lock once via ProviderLocks::for_domain instead of repeating the lock acquisition in both dispatch matches, share the IpQualitySync payload load between the single-agent and fan-out paths, and collapse the hand-rolled first-error folds to get_or_insert/map_or. Document that the capability-change handler's full-domain reconcile is deliberate, and drop leftover doubled blank lines in the mutation routers. --- crates/server/src/router/api/ip_quality.rs | 1 - crates/server/src/router/api/network_probe.rs | 1 - crates/server/src/router/api/ping.rs | 1 - crates/server/src/router/ws/agent/security.rs | 4 + crates/server/src/service/agent_reconcile.rs | 86 ++++++++----------- 5 files changed, 38 insertions(+), 55 deletions(-) diff --git a/crates/server/src/router/api/ip_quality.rs b/crates/server/src/router/api/ip_quality.rs index 2a8b9c1cd..b2ef19cdc 100644 --- a/crates/server/src/router/api/ip_quality.rs +++ b/crates/server/src/router/api/ip_quality.rs @@ -248,7 +248,6 @@ async fn update_settings( ok(setting) } - #[utoipa::path( post, path = "/api/ip-quality/servers/{id}/check", diff --git a/crates/server/src/router/api/network_probe.rs b/crates/server/src/router/api/network_probe.rs index c1493713c..eff3f9128 100644 --- a/crates/server/src/router/api/network_probe.rs +++ b/crates/server/src/router/api/network_probe.rs @@ -183,7 +183,6 @@ async fn update_setting( ok(input) } - // --------------------------------------------------------------------------- // Per-server read handlers (mounted in server.rs) // --------------------------------------------------------------------------- diff --git a/crates/server/src/router/api/ping.rs b/crates/server/src/router/api/ping.rs index 9e2ed4c65..9119435c8 100644 --- a/crates/server/src/router/api/ping.rs +++ b/crates/server/src/router/api/ping.rs @@ -139,7 +139,6 @@ async fn delete_task( ok("ok") } - #[derive(Deserialize, utoipa::IntoParams)] pub struct RecordsQuery { from: DateTime, diff --git a/crates/server/src/router/ws/agent/security.rs b/crates/server/src/router/ws/agent/security.rs index 3b0dbe73c..dcd29c594 100644 --- a/crates/server/src/router/ws/agent/security.rs +++ b/crates/server/src/router/ws/agent/security.rs @@ -186,6 +186,10 @@ pub(super) async fn on_capabilities_changed( tracing::error!("Failed to mirror capabilities for {server_id}: {e}"); } + // Deliberately a full-domain reconcile even though only the ping and + // firewall projections read capabilities: capability changes are rare, + // and "capability change ⇒ every domain converges" is a simpler invariant + // to trust than tracking which projections are capability-sensitive. if let Err(error) = state .agent_desired_state .reconcile_connection(server_id) diff --git a/crates/server/src/service/agent_reconcile.rs b/crates/server/src/service/agent_reconcile.rs index cd6b23ae1..45ab5455b 100644 --- a/crates/server/src/service/agent_reconcile.rs +++ b/crates/server/src/service/agent_reconcile.rs @@ -15,7 +15,7 @@ use serverbee_common::constants::{ CAP_DEFAULT, CAP_FIREWALL_BLOCK, has_capability, probe_type_to_cap, }; use serverbee_common::firewall::FIREWALL_MIN_PROTOCOL; -use serverbee_common::protocol::ServerMessage; +use serverbee_common::protocol::{ServerMessage, UnlockServiceDef}; use serverbee_common::types::{NetworkProbeTarget, PingTaskConfig}; use tokio::sync::Mutex; @@ -49,6 +49,17 @@ struct ProviderLocks { firewall: Mutex<()>, } +impl ProviderLocks { + fn for_domain(&self, domain: AgentDesiredStateDomain) -> &Mutex<()> { + match domain { + AgentDesiredStateDomain::PingTasks => &self.ping_tasks, + AgentDesiredStateDomain::NetworkProbes => &self.network_probes, + AgentDesiredStateDomain::IpQuality => &self.ip_quality, + AgentDesiredStateDomain::Firewall => &self.firewall, + } + } +} + /// Owns projection of persisted desired state into agent protocol messages. #[derive(Clone)] pub struct AgentDesiredStateReconciler { @@ -90,16 +101,11 @@ impl AgentDesiredStateReconciler { error = %error, "failed to reconcile agent desired state" ); - if first_error.is_none() { - first_error = Some(error); - } + first_error.get_or_insert(error); } } - match first_error { - Some(error) => Err(error), - None => Ok(()), - } + first_error.map_or(Ok(()), Err) } /// Reconcile one provider for one agent. @@ -108,27 +114,20 @@ impl AgentDesiredStateReconciler { server_id: &str, domain: AgentDesiredStateDomain, ) -> Result<(), AppError> { + let _guard = self.locks.for_domain(domain).lock().await; match domain { AgentDesiredStateDomain::PingTasks => { - let _guard = self.locks.ping_tasks.lock().await; let tasks = self.load_ping_tasks().await?; let capabilities = self.capabilities_for_agent(server_id).await?; self.send_ping_tasks(server_id, &tasks, capabilities).await; Ok(()) } AgentDesiredStateDomain::NetworkProbes => { - let _guard = self.locks.network_probes.lock().await; let setting = NetworkProbeService::get_setting(&self.db).await?; self.send_network_probes(server_id, &setting).await } - AgentDesiredStateDomain::IpQuality => { - let _guard = self.locks.ip_quality.lock().await; - self.send_ip_quality(server_id).await - } - AgentDesiredStateDomain::Firewall => { - let _guard = self.locks.firewall.lock().await; - self.send_firewall(server_id).await - } + AgentDesiredStateDomain::IpQuality => self.send_ip_quality(server_id).await, + AgentDesiredStateDomain::Firewall => self.send_firewall(server_id).await, } } @@ -139,23 +138,14 @@ impl AgentDesiredStateReconciler { &self, domain: AgentDesiredStateDomain, ) -> Result<(), AppError> { + let _guard = self.locks.for_domain(domain).lock().await; match domain { - AgentDesiredStateDomain::PingTasks => { - let _guard = self.locks.ping_tasks.lock().await; - self.reconcile_connected_ping_tasks().await - } + AgentDesiredStateDomain::PingTasks => self.reconcile_connected_ping_tasks().await, AgentDesiredStateDomain::NetworkProbes => { - let _guard = self.locks.network_probes.lock().await; self.reconcile_connected_network_probes().await } - AgentDesiredStateDomain::IpQuality => { - let _guard = self.locks.ip_quality.lock().await; - self.reconcile_connected_ip_quality().await - } - AgentDesiredStateDomain::Firewall => { - let _guard = self.locks.firewall.lock().await; - self.reconcile_connected_firewall().await - } + AgentDesiredStateDomain::IpQuality => self.reconcile_connected_ip_quality().await, + AgentDesiredStateDomain::Firewall => self.reconcile_connected_firewall().await, } } @@ -210,26 +200,16 @@ impl AgentDesiredStateReconciler { error = %error, "failed to reconcile network probe desired state" ); - if first_error.is_none() { - first_error = Some(error); - } + first_error.get_or_insert(error); } } - match first_error { - Some(error) => Err(error), - None => Ok(()), - } + first_error.map_or(Ok(()), Err) } async fn reconcile_connected_ip_quality(&self) -> Result<(), AppError> { let server_ids = self.agent_manager.connected_server_ids(); - let services = IpQualityService::enabled_service_defs(&self.db).await?; - let interval_hours = ip_quality_interval_hours( - IpQualityService::get_setting(&self.db) - .await? - .check_interval_hours, - )?; + let (services, interval_hours) = self.load_ip_quality().await?; for server_id in server_ids { self.send_if_online( @@ -255,16 +235,11 @@ impl AgentDesiredStateReconciler { error = %error, "failed to reconcile firewall desired state" ); - if first_error.is_none() { - first_error = Some(error); - } + first_error.get_or_insert(error); } } - match first_error { - Some(error) => Err(error), - None => Ok(()), - } + first_error.map_or(Ok(()), Err) } async fn load_ping_tasks(&self) -> Result, AppError> { @@ -360,13 +335,20 @@ impl AgentDesiredStateReconciler { Ok(()) } - async fn send_ip_quality(&self, server_id: &str) -> Result<(), AppError> { + /// Load the enabled unlock-service catalog and check interval that every + /// `IpQualitySync` frame carries. + async fn load_ip_quality(&self) -> Result<(Vec, u32), AppError> { let services = IpQualityService::enabled_service_defs(&self.db).await?; let interval_hours = ip_quality_interval_hours( IpQualityService::get_setting(&self.db) .await? .check_interval_hours, )?; + Ok((services, interval_hours)) + } + + async fn send_ip_quality(&self, server_id: &str) -> Result<(), AppError> { + let (services, interval_hours) = self.load_ip_quality().await?; self.send_if_online( server_id, ServerMessage::IpQualitySync { From ab1b88abc2c6abff53943f601b37f12918379b44 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 15:53:24 +0800 Subject: [PATCH 15/44] refactor(agent): narrow ConnectionRuntime field visibility Only terminal_manager, unlock_checker and cmd_result_tx are touched by the reporter select! loop; the remaining fields are runtime-internal. --- crates/agent/src/reporter/runtime.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/crates/agent/src/reporter/runtime.rs b/crates/agent/src/reporter/runtime.rs index a7be69d85..9b19db0d1 100644 --- a/crates/agent/src/reporter/runtime.rs +++ b/crates/agent/src/reporter/runtime.rs @@ -58,18 +58,18 @@ pub(super) struct ConnectionEvents { /// Everything a single agent connection needs to execute server commands. pub(super) struct ConnectionRuntime { - pub(super) capabilities: Arc, - pub(super) ping_manager: PingManager, + capabilities: Arc, + ping_manager: PingManager, pub(super) terminal_manager: TerminalManager, - pub(super) network_prober: NetworkProber, - pub(super) file_manager: FileManager, + network_prober: NetworkProber, + file_manager: FileManager, pub(super) unlock_checker: UnlockChecker, pub(super) cmd_result_tx: mpsc::Sender, file_tx: mpsc::Sender, docker_tx: mpsc::Sender, - pub(super) docker_manager: Option, + docker_manager: Option, docker_available: bool, - pub(super) docker_stats_interval: Option, + docker_stats_interval: Option, docker_retry_interval: tokio::time::Interval, upgrade_cfg: UpgradeConfig, firewall_manager: Arc, From a6c67023cab3f1fae5d9b3162a71ade1e9be539e Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 15:53:35 +0800 Subject: [PATCH 16/44] test(server): share the first-connect noise predicate across suites MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ten integration suites carried their own copy of the ignorable first-connect push list, and each new sync frame (latest: ip_quality_sync) meant patching every copy — one had already drifted. Hoist the predicate into tests/common and replace integration.rs's drain_through_ack frame counts with a drain-until-quiet loop so connect-flow changes stop breaking unrelated suites. --- crates/server/tests/agent_docker_extra.rs | 17 +------ crates/server/tests/agent_file_ops.rs | 22 ++------- crates/server/tests/agent_ws_dispatch.rs | 29 +++--------- crates/server/tests/common/mod.rs | 19 ++++++++ crates/server/tests/docker_integration.rs | 19 +++----- crates/server/tests/integration.rs | 49 ++++++++++----------- crates/server/tests/router_admin_extra.rs | 9 +--- crates/server/tests/router_file_extra.rs | 31 ++++--------- crates/server/tests/router_probes_extra.rs | 4 +- crates/server/tests/router_server_extra2.rs | 11 +---- crates/server/tests/ws_docker_browser.rs | 19 ++------ crates/server/tests/ws_terminal_relay.rs | 19 ++------ 12 files changed, 80 insertions(+), 168 deletions(-) diff --git a/crates/server/tests/agent_docker_extra.rs b/crates/server/tests/agent_docker_extra.rs index 7563713f9..9c9995cc5 100644 --- a/crates/server/tests/agent_docker_extra.rs +++ b/crates/server/tests/agent_docker_extra.rs @@ -18,6 +18,7 @@ mod common; use common::{ + is_first_connect_noise, connect_agent, http_client, login_admin, login_as_new_user, recv_agent_text, register_agent, start_test_server, AgentReader, AgentSink, }; @@ -68,22 +69,6 @@ async fn send_docker_system_info(sink: &mut AgentSink, reader: &mut AgentReader, } } -/// Drain the first-connect pushes the server emits to a default agent. These -/// (ping/network sync + firewall blocklist messages) are unrelated to the -/// docker flow under test, so they must be ignored by every agent responder. -fn is_first_connect_noise(msg_type: Option<&str>) -> bool { - matches!( - msg_type, - Some("ping_tasks_sync") - | Some("network_probe_sync") - | Some("ip_quality_sync") - | Some("blocklist_reset") - | Some("blocklist_sync") - | Some("blocklist_add") - | Some("blocklist_remove") - ) -} - // --------------------------------------------------------------------------- // GET /servers/{id}/docker/containers (cache read; no agent responder needed) // --------------------------------------------------------------------------- diff --git a/crates/server/tests/agent_file_ops.rs b/crates/server/tests/agent_file_ops.rs index f48b05a57..81a4ede4f 100644 --- a/crates/server/tests/agent_file_ops.rs +++ b/crates/server/tests/agent_file_ops.rs @@ -15,6 +15,7 @@ mod common; use common::{ + is_first_connect_noise, AgentReader, AgentSink, connect_agent, http_client, login_admin, login_as_new_user, recv_agent_text, register_agent, send_system_info, start_test_server, }; @@ -40,21 +41,6 @@ async fn connect_agent_with_caps( (sink, reader) } -/// First-connect pushes a default agent (which reports CAP_FIREWALL_BLOCK) must -/// tolerate and ignore. These are unrelated to the file flow under test. -fn is_ignorable_push(msg_type: Option<&str>) -> bool { - matches!( - msg_type, - Some("ping_tasks_sync") - | Some("network_probe_sync") - | Some("ip_quality_sync") - | Some("blocklist_reset") - | Some("blocklist_sync") - | Some("blocklist_add") - | Some("blocklist_remove") - ) -} - /// Spawn a responder that waits for a single forwarded request whose `type` /// equals `request_type`, then sends `build_response(msg_id)` back to the /// server. Ignorable first-connect pushes are skipped; any other command is a @@ -82,7 +68,7 @@ where .expect("send file-op response"); return; } - other if is_ignorable_push(other) => {} + other if is_first_connect_noise(other) => {} Some(other) => panic!("unexpected agent command: {other}"), None => {} } @@ -648,7 +634,7 @@ async fn test_start_download_happy_path() { assert_eq!(msg["path"], "/var/log/syslog"); return msg["transfer_id"].as_str().map(str::to_string); } - other if is_ignorable_push(other) => {} + other if is_first_connect_noise(other) => {} Some(other) => panic!("unexpected agent command: {other}"), None => {} } @@ -786,7 +772,7 @@ async fn test_upload_file_happy_path() { .expect("send upload_complete"); return; } - other if is_ignorable_push(other) => {} + other if is_first_connect_noise(other) => {} Some(other) => panic!("unexpected agent command: {other}"), None => {} } diff --git a/crates/server/tests/agent_ws_dispatch.rs b/crates/server/tests/agent_ws_dispatch.rs index 999cc93a3..ace2ca7f3 100644 --- a/crates/server/tests/agent_ws_dispatch.rs +++ b/crates/server/tests/agent_ws_dispatch.rs @@ -40,6 +40,7 @@ mod common; use common::{ + is_first_connect_noise, connect_agent, http_client, login_admin, login_as_new_user, recv_agent_text, register_agent, send_system_info, start_test_server, AgentReader, AgentSink, }; @@ -60,22 +61,6 @@ async fn send_agent_frame(sink: &mut AgentSink, frame: Value) { .expect("send agent frame"); } -/// First-connect pushes that a default agent must tolerate and ignore. These -/// (ping/network/IP-quality sync + firewall blocklist messages) are unrelated to -/// the dispatch arms under test. -fn is_ignorable_push(msg_type: Option<&str>) -> bool { - matches!( - msg_type, - Some("ping_tasks_sync") - | Some("network_probe_sync") - | Some("ip_quality_sync") - | Some("blocklist_reset") - | Some("blocklist_sync") - | Some("blocklist_add") - | Some("blocklist_remove") - ) -} - /// Drain any frames the server pushes right after the SystemInfo handshake. /// Returns once the inbound stream is quiet for `quiet_ms`. async fn drain_first_connect_pushes(reader: &mut AgentReader, quiet_ms: u64) { @@ -206,7 +191,7 @@ async fn test_traceroute_round_update_completed_persists_and_is_queryable() { send_agent_frame(&mut sink, reply).await; return; } - other if is_ignorable_push(other) => {} + other if is_first_connect_noise(other) => {} Some(_other) => {} None => {} } @@ -305,7 +290,7 @@ async fn test_traceroute_round_update_progression_then_complete() { .await; return; } - other if is_ignorable_push(other) => {} + other if is_first_connect_noise(other) => {} Some(_other) => {} None => {} } @@ -507,7 +492,7 @@ async fn test_legacy_traceroute_result_records_legacy_protocol() { .await; return; } - other if is_ignorable_push(other) => {} + other if is_first_connect_noise(other) => {} Some(_other) => {} None => {} } @@ -595,7 +580,7 @@ async fn test_upgrade_progress_is_acked() { assert_eq!(ack["msg_id"], "up-prog-1"); return; } - other if is_ignorable_push(other) => {} + other if is_first_connect_noise(other) => {} Some(_other) => {} None => {} } @@ -651,7 +636,7 @@ async fn test_upgrade_result_acks_and_clears_running_job() { // agent ONLINE. The server's Ack for this result arrives // as a later frame the loop simply ignores. } - other if is_ignorable_push(other) => {} + other if is_first_connect_noise(other) => {} Some(_other) => {} None => {} } @@ -724,7 +709,7 @@ async fn test_capability_denied_upgrade_marks_job_failed() { // (returning here would drop the socket → agent offline → // the retrigger 404s instead of returning the expected 200). } - other if is_ignorable_push(other) => {} + other if is_first_connect_noise(other) => {} Some(_other) => {} None => {} } diff --git a/crates/server/tests/common/mod.rs b/crates/server/tests/common/mod.rs index f68c98377..c268144b6 100644 --- a/crates/server/tests/common/mod.rs +++ b/crates/server/tests/common/mod.rs @@ -275,3 +275,22 @@ pub async fn send_system_info( } } } + +/// First-connect pushes the server emits to a default agent right after the +/// SystemInfo handshake: the desired-state syncs (ping tasks, network probes, +/// IP quality) plus the firewall blocklist frames. Agent responders in tests +/// must tolerate and ignore all of them — they are unrelated to whatever flow +/// a given test exercises, and this single predicate is the one place to +/// extend when a new first-connect sync frame is introduced. +pub fn is_first_connect_noise(msg_type: Option<&str>) -> bool { + matches!( + msg_type, + Some("ping_tasks_sync") + | Some("network_probe_sync") + | Some("ip_quality_sync") + | Some("blocklist_reset") + | Some("blocklist_sync") + | Some("blocklist_add") + | Some("blocklist_remove") + ) +} diff --git a/crates/server/tests/docker_integration.rs b/crates/server/tests/docker_integration.rs index 583a59334..0aa028d9f 100644 --- a/crates/server/tests/docker_integration.rs +++ b/crates/server/tests/docker_integration.rs @@ -15,6 +15,9 @@ use tokio_tungstenite::tungstenite; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_tungstenite::tungstenite::http::HeaderValue; +mod common; +use common::is_first_connect_noise; + async fn start_test_server() -> (String, tempfile::TempDir) { let tmp = tempfile::tempdir().expect("Failed to create temp dir"); let data_dir = tmp.path().to_str().unwrap().to_string(); @@ -409,13 +412,7 @@ async fn test_docker_info_endpoint_requests_agent_when_cache_empty() { // A default agent reports CAP_FIREWALL_BLOCK, so the first-connect // path also pushes blocklist reset/sync. These are unrelated to the // docker flow under test and are ignored. - Some("ping_tasks_sync") - | Some("network_probe_sync") - | Some("ip_quality_sync") - | Some("blocklist_reset") - | Some("blocklist_sync") - | Some("blocklist_add") - | Some("blocklist_remove") => {} + noise if is_first_connect_noise(noise) => {} Some(other) => panic!("Unexpected agent command: {other}"), None => {} } @@ -787,13 +784,7 @@ async fn test_docker_unavailable_fails_pending_request_and_clears_feature_state( // A default agent reports CAP_FIREWALL_BLOCK, so the first-connect // path also pushes blocklist reset/sync. These are unrelated to the // docker flow under test and are ignored. - Some("ping_tasks_sync") - | Some("network_probe_sync") - | Some("ip_quality_sync") - | Some("blocklist_reset") - | Some("blocklist_sync") - | Some("blocklist_add") - | Some("blocklist_remove") => {} + noise if is_first_connect_noise(noise) => {} Some(other) => panic!("Unexpected agent command: {other}"), None => {} } diff --git a/crates/server/tests/integration.rs b/crates/server/tests/integration.rs index d602f8f80..3bf7e6cc4 100644 --- a/crates/server/tests/integration.rs +++ b/crates/server/tests/integration.rs @@ -4358,9 +4358,12 @@ mod firewall_tests { } /// Drain messages from `reader` until a `system_info` `Ack` for `msg_id` - /// is observed. Returns every non-ack message read along the way (in - /// order) so the caller can inspect post-connect server pushes such as - /// `BlocklistReset` / `BlocklistSync`. + /// is observed, then keep draining until the socket stays quiet for a + /// beat (the post-connect reconcile pushes arrive back-to-back). Returns + /// every non-ack message read along the way (in order) so the caller can + /// inspect post-connect server pushes such as `BlocklistReset` / + /// `BlocklistSync`, without hard-coding how many sync frames a connect + /// currently produces. async fn drain_through_ack( reader: &mut futures_util::stream::SplitStream< tokio_tungstenite::WebSocketStream< @@ -4368,17 +4371,22 @@ mod firewall_tests { >, >, msg_id: &str, - extra_after_ack: usize, ) -> Vec { let mut collected: Vec = Vec::new(); let mut acked = false; - let mut extra_remaining = extra_after_ack; loop { - let msg = tokio::time::timeout(Duration::from_secs(3), reader.next()) - .await - .expect("timeout waiting for ws message") - .expect("ws closed") - .expect("ws read error"); + let timeout = if acked { + // Post-ack pushes are emitted by the same handler in one + // burst; a quiet gap this long means the connect flow is done. + Duration::from_millis(800) + } else { + Duration::from_secs(3) + }; + let msg = match tokio::time::timeout(timeout, reader.next()).await { + Ok(msg) => msg.expect("ws closed").expect("ws read error"), + Err(_) if acked => break, + Err(_) => panic!("timeout waiting for ws message"), + }; let text = match msg { tungstenite::Message::Text(t) => t.to_string(), tungstenite::Message::Ping(_) | tungstenite::Message::Pong(_) => continue, @@ -4387,18 +4395,9 @@ mod firewall_tests { let parsed: serde_json::Value = serde_json::from_str(&text).expect("parse json"); if parsed["type"] == "ack" && parsed["msg_id"] == msg_id { acked = true; - if extra_remaining == 0 { - break; - } continue; } collected.push(parsed); - if acked { - extra_remaining -= 1; - if extra_remaining == 0 { - break; - } - } } collected } @@ -4451,9 +4450,9 @@ mod firewall_tests { serverbee_common::constants::PROTOCOL_VERSION, ) .await; - // AgentReady reconciliation follows the ack: ping, network, IP quality, - // then BlocklistReset+BlocklistSync. Drain all five messages. - let _ = drain_through_ack(&mut reader, "pb-1", 5).await; + // AgentReady reconciliation pushes follow the ack (ping/network/IP + // quality sync, BlocklistReset+BlocklistSync); drain them all. + let _ = drain_through_ack(&mut reader, "pb-1").await; // Now POST a block and assert the agent receives BlocklistAdd. let resp = admin @@ -4572,7 +4571,7 @@ mod firewall_tests { ) .await; - let extras = drain_through_ack(&mut reader, "connect-sync", 5).await; + let extras = drain_through_ack(&mut reader, "connect-sync").await; // Reset, then Sync are appended after the ack handler runs. let types: Vec<&str> = extras .iter() @@ -4609,7 +4608,7 @@ mod firewall_tests { serverbee_common::constants::PROTOCOL_VERSION, ) .await; - let _drain = drain_through_ack(&mut reader, "ack-1", 5).await; + let _drain = drain_through_ack(&mut reader, "ack-1").await; let resp = admin .post(format!("{}/api/firewall/blocks", base_url)) @@ -4683,7 +4682,7 @@ mod firewall_tests { 1, ) .await; - let extras = drain_through_ack(&mut reader, "old-1", 3).await; + let extras = drain_through_ack(&mut reader, "old-1").await; let types: Vec<&str> = extras .iter() .map(|m| m["type"].as_str().unwrap_or("")) diff --git a/crates/server/tests/router_admin_extra.rs b/crates/server/tests/router_admin_extra.rs index e43b97ca7..6462cac15 100644 --- a/crates/server/tests/router_admin_extra.rs +++ b/crates/server/tests/router_admin_extra.rs @@ -153,13 +153,8 @@ async fn run_scheduled_task_dispatches_to_agent_and_persists_result() { .await; return; } - // First-connect noise from a default agent: ignore. - Some("ping_tasks_sync") - | Some("network_probe_sync") - | Some("blocklist_reset") - | Some("blocklist_sync") - | Some("blocklist_add") - | Some("blocklist_remove") => {} + // First-connect noise from a default agent (and anything else + // unrelated to the scheduled run): ignore. _ => {} } } diff --git a/crates/server/tests/router_file_extra.rs b/crates/server/tests/router_file_extra.rs index edb52469b..88c3f30d9 100644 --- a/crates/server/tests/router_file_extra.rs +++ b/crates/server/tests/router_file_extra.rs @@ -22,6 +22,7 @@ mod common; use common::{ + is_first_connect_noise, AgentReader, AgentSink, connect_agent, http_client, login_admin, login_as_new_user, recv_agent_text, register_agent, send_system_info, start_test_server, }; @@ -46,22 +47,6 @@ async fn connect_agent_with_caps( (sink, reader) } -/// First-connect pushes the server emits to a default agent that the file flow -/// has no interest in. Ignore them so protocol drift on the real commands still -/// surfaces loudly. -fn is_ignorable_push(msg_type: Option<&str>) -> bool { - matches!( - msg_type, - Some("ping_tasks_sync") - | Some("network_probe_sync") - | Some("ip_quality_sync") - | Some("blocklist_reset") - | Some("blocklist_sync") - | Some("blocklist_add") - | Some("blocklist_remove") - ) -} - /// Spawn a responder that waits for one forwarded request whose `type` equals /// `request_type`, then sends `build_response(msg_id)` back to the server. fn spawn_single_response( @@ -87,7 +72,7 @@ where .expect("send file-op response"); return; } - other if is_ignorable_push(other) => {} + other if is_first_connect_noise(other) => {} Some(other) => panic!("unexpected agent command: {other}"), None => {} } @@ -411,7 +396,7 @@ async fn test_delete_dir_recursive_happy_path() { .expect("send file_op_result"); return; } - other if is_ignorable_push(other) => {} + other if is_first_connect_noise(other) => {} Some(other) => panic!("unexpected agent command: {other}"), None => {} } @@ -577,7 +562,7 @@ async fn test_download_full_lifecycle_streams_bytes() { } return; } - other if is_ignorable_push(other) => {} + other if is_first_connect_noise(other) => {} Some(other) => panic!("unexpected agent command: {other}"), None => {} } @@ -661,7 +646,7 @@ async fn test_download_by_id_not_ready_is_400() { let msg = recv_agent_text(&mut reader).await; match msg["type"].as_str() { Some("file_download_start") => return, - other if is_ignorable_push(other) => {} + other if is_first_connect_noise(other) => {} Some(other) => panic!("unexpected agent command: {other}"), None => {} } @@ -709,7 +694,7 @@ async fn test_download_by_id_cross_user_is_404() { let msg = recv_agent_text(&mut reader).await; match msg["type"].as_str() { Some("file_download_start") => return, - other if is_ignorable_push(other) => {} + other if is_first_connect_noise(other) => {} Some(other) => panic!("unexpected agent command: {other}"), None => {} } @@ -779,7 +764,7 @@ async fn test_list_then_cancel_owned_download_transfer() { } return; } - other if is_ignorable_push(other) => {} + other if is_first_connect_noise(other) => {} Some(other) => panic!("unexpected agent command: {other}"), None => {} } @@ -944,7 +929,7 @@ async fn test_upload_agent_rejects_start_is_400() { .expect("send upload_error"); return; } - other if is_ignorable_push(other) => {} + other if is_first_connect_noise(other) => {} Some(other) => panic!("unexpected agent command: {other}"), None => {} } diff --git a/crates/server/tests/router_probes_extra.rs b/crates/server/tests/router_probes_extra.rs index d6ce85c0c..761791ae2 100644 --- a/crates/server/tests/router_probes_extra.rs +++ b/crates/server/tests/router_probes_extra.rs @@ -12,6 +12,7 @@ mod common; use common::{ + is_first_connect_noise, AgentReader, AgentSink, connect_agent, create_server, http_client, login_admin, login_as_new_user, recv_agent_text, register_agent, send_system_info, start_test_server, }; @@ -50,8 +51,7 @@ async fn recv_until(reader: &mut AgentReader, expected: &str) -> Value { match ty { // First-connect noise a default agent always receives; unrelated to // the request under test, so it is ignored. - "ping_tasks_sync" | "network_probe_sync" | "blocklist_reset" | "blocklist_sync" - | "blocklist_add" | "blocklist_remove" | "ip_quality_sync" => {} + noise if is_first_connect_noise(Some(noise)) => {} other => panic!("unexpected agent command while awaiting `{expected}`: {other}"), } } diff --git a/crates/server/tests/router_server_extra2.rs b/crates/server/tests/router_server_extra2.rs index 03923cb35..226d516ae 100644 --- a/crates/server/tests/router_server_extra2.rs +++ b/crates/server/tests/router_server_extra2.rs @@ -36,6 +36,7 @@ mod common; use common::{ + is_first_connect_noise, connect_agent, create_server, http_client, login_admin, login_as_new_user, recv_agent_text, register_agent, send_system_info, start_test_server, AgentSink, AgentReader, }; @@ -69,15 +70,7 @@ async fn online_agent(base_url: &str, token: &str, caps: u32) -> (AgentSink, Age Ok(msg) => { let ty = msg["type"].as_str().unwrap_or_default(); // Ignore the expected first-connect control-plane sync frames. - if matches!( - ty, - "ping_tasks_sync" - | "network_probe_sync" - | "blocklist_reset" - | "blocklist_sync" - | "blocklist_add" - | "blocklist_remove" - ) { + if is_first_connect_noise(Some(ty)) { continue; } // Anything else is unexpected this early; keep draining anyway. diff --git a/crates/server/tests/ws_docker_browser.rs b/crates/server/tests/ws_docker_browser.rs index 27e787e2b..056f2c14c 100644 --- a/crates/server/tests/ws_docker_browser.rs +++ b/crates/server/tests/ws_docker_browser.rs @@ -31,6 +31,7 @@ mod common; use common::{ + is_first_connect_noise, connect_agent, http_client, login_admin, login_as_new_user, recv_agent_text, register_agent, start_test_server, AgentReader, AgentSink, }; @@ -79,20 +80,6 @@ async fn send_agent_frame(sink: &mut AgentSink, frame: Value) { .expect("send agent frame"); } -/// First-connect pushes a default agent must tolerate and ignore. -fn is_ignorable_push(msg_type: Option<&str>) -> bool { - matches!( - msg_type, - Some("ping_tasks_sync") - | Some("network_probe_sync") - | Some("ip_quality_sync") - | Some("blocklist_reset") - | Some("blocklist_sync") - | Some("blocklist_add") - | Some("blocklist_remove") - ) -} - /// Drain frames the server pushes right after the SystemInfo handshake; returns /// once the inbound stream is quiet for `quiet_ms`. async fn drain_first_connect_pushes(reader: &mut AgentReader, quiet_ms: u64) { @@ -226,7 +213,7 @@ async fn agent_recv_until_types( let Some(msg_type) = parsed["type"].as_str() else { continue; }; - if is_ignorable_push(Some(msg_type)) { + if is_first_connect_noise(Some(msg_type)) { continue; } if expected.contains(&msg_type) && !seen.iter().any(|s| s == msg_type) { @@ -291,7 +278,7 @@ async fn docker_logs_relay_subscribe_logs_and_unsubscribe() { let _ = tx.send(msg.clone()).await; return; } - other if is_ignorable_push(other) => {} + other if is_first_connect_noise(other) => {} Some(_other) => {} None => {} } diff --git a/crates/server/tests/ws_terminal_relay.rs b/crates/server/tests/ws_terminal_relay.rs index 5472d618e..4fb08c1de 100644 --- a/crates/server/tests/ws_terminal_relay.rs +++ b/crates/server/tests/ws_terminal_relay.rs @@ -46,6 +46,7 @@ mod common; use common::{ + is_first_connect_noise, connect_agent, http_client, login_admin, recv_agent_text, register_agent, send_system_info, start_test_server, AgentReader, AgentSink, }; @@ -145,20 +146,6 @@ async fn send_agent_frame(sink: &mut AgentSink, frame: Value) { .expect("send agent frame"); } -/// First-connect pushes a default/terminal agent must tolerate and ignore. -fn is_ignorable_push(msg_type: Option<&str>) -> bool { - matches!( - msg_type, - Some("ping_tasks_sync") - | Some("network_probe_sync") - | Some("ip_quality_sync") - | Some("blocklist_reset") - | Some("blocklist_sync") - | Some("blocklist_add") - | Some("blocklist_remove") - ) -} - /// Drain frames the server pushes right after the SystemInfo handshake, until /// the inbound stream is quiet for `quiet_ms`. async fn drain_first_connect_pushes(reader: &mut AgentReader, quiet_ms: u64) { @@ -245,7 +232,7 @@ async fn terminal_ws_relays_session_started_output_and_input_resize() { Some("terminal_close") => { return (session_id, got_input, got_resize); } - other if is_ignorable_push(other) => {} + other if is_first_connect_noise(other) => {} _ => {} } } @@ -349,7 +336,7 @@ async fn terminal_ws_relays_agent_error_to_browser() { .await; } Some("terminal_close") => return, - other if is_ignorable_push(other) => {} + other if is_first_connect_noise(other) => {} _ => {} } } From 1d24c225adf297ab42c3848655f5f77fdc1ed98b Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 15:53:35 +0800 Subject: [PATCH 17/44] refactor(web): type detail range keys as a RangeKey union MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Derive RangeKey from the range catalog, validate the range search param on parse instead of casting it, and narrow onRangeChange callbacks so invalid keys can no longer flow into detail navigation. The raw catalogs lose their exports — rangesForVariant is the only consumer. --- .../src/components/network/network-tab.tsx | 5 ++-- .../status/server-detail-content.tsx | 6 ++--- apps/web/src/lib/server-detail-nav.ts | 27 ++++++++++++------- 3 files changed, 23 insertions(+), 15 deletions(-) diff --git a/apps/web/src/components/network/network-tab.tsx b/apps/web/src/components/network/network-tab.tsx index 291945284..1aba17964 100644 --- a/apps/web/src/components/network/network-tab.tsx +++ b/apps/web/src/components/network/network-tab.tsx @@ -19,6 +19,7 @@ import { CHART_COLORS } from '@/lib/chart-colors' import { formatDateTime } from '@/lib/format' import { getNetworkTargetDisplayName } from '@/lib/network-i18n' import type { NetworkProbeAnomaly, NetworkTargetSummary } from '@/lib/network-types' +import type { RangeKey } from '@/lib/server-detail-nav' import { cn } from '@/lib/utils' // The network tab shares the server-detail `range` search param, which uses @@ -35,7 +36,7 @@ const RANGE_KEYS = ['realtime', '1h', '6h', '24h', '7d', '30d'] as const export interface NetworkTabProps { /** Called when the viewer picks a new range; keys are metrics-style. */ - onRangeChange?: (rangeKey: string) => void + onRangeChange?: (rangeKey: RangeKey) => void /** Metrics-style range key from the shared `range` search param. */ rangeKey?: string serverId: string @@ -86,7 +87,7 @@ function TimeRangeControls({ rangeKey, t }: { - onChange: (rangeKey: string) => void + onChange: (rangeKey: RangeKey) => void rangeKey: string t: (key: string, opts?: Record) => string }) { diff --git a/apps/web/src/components/status/server-detail-content.tsx b/apps/web/src/components/status/server-detail-content.tsx index 1f57f358f..31067458c 100644 --- a/apps/web/src/components/status/server-detail-content.tsx +++ b/apps/web/src/components/status/server-detail-content.tsx @@ -26,7 +26,7 @@ import type { } from '@/lib/api-schema' import { buildMergedDiskIoSeries, buildPerDiskIoSeries } from '@/lib/disk-io' import { type ServerMetrics, useLiveServers } from '@/lib/server-catalog' -import { rangesForVariant, resolveRange, type TimeRange } from '@/lib/server-detail-nav' +import { type RangeKey, rangesForVariant, resolveRange, type TimeRange } from '@/lib/server-detail-nav' import { cn, formatBytes } from '@/lib/utils' import { computeAggregateUptime } from '@/lib/widget-helpers' @@ -48,7 +48,7 @@ export interface ServerDetailContentProps { * page `show_network` toggle). */ networkTab?: React.ReactNode /** Called by range buttons when the viewer picks a new historical window. */ - onRangeChange?: (rangeKey: string) => void + onRangeChange?: (rangeKey: RangeKey) => void /** Called when the viewer switches detail tabs. */ onTabChange?: (tab: string) => void /** Currently selected range key from the URL or local state. */ @@ -539,7 +539,7 @@ function MetricsTabContent({ name: string }[] gpuChartData: Record[] - onRangeChange?: (rangeKey: string) => void + onRangeChange?: (rangeKey: RangeKey) => void rangeIndex: number ranges: TimeRange[] formatTime: ((time: string) => string) | undefined diff --git a/apps/web/src/lib/server-detail-nav.ts b/apps/web/src/lib/server-detail-nav.ts index 05746694c..ef002fe4f 100644 --- a/apps/web/src/lib/server-detail-nav.ts +++ b/apps/web/src/lib/server-detail-nav.ts @@ -19,20 +19,20 @@ export type PublicServerDetailTab = (typeof PUBLIC_SERVER_DETAIL_TABS)[number] /** Admin detail search: `range` always present (default realtime); `tab` kept * optional so old `/servers/$id` links (range only) stay valid. */ -export function parseServerDetailSearch(search: Record): { range: string; tab?: ServerDetailTab } { - const range = (search.range as string) || 'realtime' +export function parseServerDetailSearch(search: Record): { range: RangeKey; tab?: ServerDetailTab } { + const range = isRangeKey(search.range) ? search.range : 'realtime' return SERVER_DETAIL_TABS.includes(search.tab as ServerDetailTab) ? { range, tab: search.tab as ServerDetailTab } : { range } } -/** Public detail search: both params optional; unknown tabs are dropped. */ +/** Public detail search: both params optional; unknown ranges and tabs are dropped. */ export function parsePublicServerDetailSearch(search: Record): { - range?: string + range?: RangeKey tab?: PublicServerDetailTab } { return { - range: typeof search.range === 'string' ? search.range : undefined, + range: isRangeKey(search.range) ? search.range : undefined, ...(PUBLIC_SERVER_DETAIL_TABS.includes(search.tab as PublicServerDetailTab) ? { tab: search.tab as PublicServerDetailTab } : {}) @@ -50,14 +50,21 @@ export function resolvePublicActiveTab( // --- Range catalogs ---------------------------------------------------------- +const RANGE_KEYS = ['realtime', '1h', '6h', '24h', '7d', '30d'] as const +export type RangeKey = (typeof RANGE_KEYS)[number] + +function isRangeKey(value: unknown): value is RangeKey { + return (RANGE_KEYS as readonly unknown[]).includes(value) +} + export interface TimeRange { hours: number interval: string - key: string + key: RangeKey label: string } -export const ADMIN_TIME_RANGES: TimeRange[] = [ +const ADMIN_TIME_RANGES: TimeRange[] = [ { key: 'realtime', label: 'range_realtime', hours: 0, interval: 'realtime' }, { key: '1h', label: 'range_1h', hours: 1, interval: 'raw' }, { key: '6h', label: 'range_6h', hours: 6, interval: 'raw' }, @@ -69,7 +76,7 @@ export const ADMIN_TIME_RANGES: TimeRange[] = [ // Public variant cannot rely on WS-driven realtime metrics, so realtime is // dropped; everything else mirrors the admin range options because the // public metrics endpoint accepts the same `interval` query parameter. -export const PUBLIC_TIME_RANGES: TimeRange[] = ADMIN_TIME_RANGES.filter((r) => r.key !== 'realtime') +const PUBLIC_TIME_RANGES: TimeRange[] = ADMIN_TIME_RANGES.filter((r) => r.key !== 'realtime') export function rangesForVariant(variant: 'admin' | 'public'): TimeRange[] { return variant === 'public' ? PUBLIC_TIME_RANGES : ADMIN_TIME_RANGES @@ -86,8 +93,8 @@ export function resolveRange(rangeKey: string | undefined, ranges: TimeRange[]) /** The retired standalone network page counted hours in its `range` param * ('1' | '6' | ...); detail URLs use metrics-style keys. Used by the legacy * `/network/$serverId` redirect so old bookmarks land on the right window. */ -export function legacyNetworkRangeToRangeKey(range: string): string { - const mapping: Record = { +export function legacyNetworkRangeToRangeKey(range: string): RangeKey { + const mapping: Record = { realtime: 'realtime', '1': '1h', '6': '6h', From d8bbeea3321ff4bee9833123c7c1c66a6ec7f580 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 17:13:57 +0800 Subject: [PATCH 18/44] refactor: type browser update frames as a LiveMetrics partial projection BrowserMessage::Update carried a full ServerStatus with static fields zero-filled, an invariant enforced only by comments. Every client re-encoded it defensively: the web merge hand-picked live fields and iOS guarded totals with >0 checks but missed swapTotal, which updates stomped to 0 on every frame. Update now carries a dedicated LiveMetrics type holding only the fields an agent report can populate. The wire is compatible by subtraction (same tag, same keys, fewer of them), so older iOS builds decode the absent statics as nil and keep cached values. The web merge becomes a plain spread, update frames can no longer seed the catalog cache with placeholder rows, and the unused name field is dropped from updates. See docs/adr/0002-update-frames-carry-live-metrics.md. --- CONTEXT.md | 2 + apps/web/src/hooks/use-servers-ws.ts | 4 +- apps/web/src/lib/server-catalog.test.ts | 58 +++++++++----- apps/web/src/lib/server-catalog.ts | 80 ++++++++----------- crates/common/src/protocol.rs | 2 +- crates/common/src/types.rs | 34 ++++++++ crates/server/src/service/agent_manager.rs | 76 ++++++++++++------ .../0002-update-frames-carry-live-metrics.md | 26 ++++++ 8 files changed, 189 insertions(+), 93 deletions(-) create mode 100644 docs/adr/0002-update-frames-carry-live-metrics.md diff --git a/CONTEXT.md b/CONTEXT.md index e68fac099..414cbcf63 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -4,3 +4,5 @@ - **Network tab** — the server-detail tab hosting the network quality detail. Admin gets the full experience (chart, traceroute, target management, CSV export); the public status variant is the redacted summary (targets + anomalies only). - **Server detail tabs** — the top-level structure of the server detail page: Metrics (default, includes the cost/traffic/uptime overview blocks), Network, Traffic, Security, IP Quality. Tab and time range are URL-driven (`?tab=`, `?range=`); the `range` window is shared across tabs. - **Standalone public network page** — `/status/network/$serverId`; exists only as a fallback for status pages configured with `show_server_detail=false` but `show_network=true` (see ADR-0001). +- **Live metrics (实时指标帧)** — the partial projection of a server carried by WS `update` frames: only the fields an agent report can populate (usage, speeds, loads, connection counts). Static facts are absent from the wire; clients keep their cached values when merging (see ADR-0002). _Avoid_: partial status, update payload. +- **Full server status** — the complete ~35-field snapshot (`ServerStatus`) carried by `full_sync` and REST. The only sources allowed to seed the catalog or overwrite static facts (totals, os, tags, geo, enrollment). diff --git a/apps/web/src/hooks/use-servers-ws.ts b/apps/web/src/hooks/use-servers-ws.ts index 5dfdc4bdc..868fa0cfb 100644 --- a/apps/web/src/hooks/use-servers-ws.ts +++ b/apps/web/src/hooks/use-servers-ws.ts @@ -6,7 +6,7 @@ import { toast } from 'sonner' import type { SecurityEventDto, SecurityEventList } from '@/lib/api-schema' import type { IpQualitySnapshotData, ServerIpQualityData, UnlockResultDto, UnlockStatus } from '@/lib/ip-quality-types' import type { NetworkProbeResultData } from '@/lib/network-types' -import { projectServerCatalog, type ServerMetrics } from '@/lib/server-catalog' +import { type LiveMetrics, projectServerCatalog, type ServerMetrics } from '@/lib/server-catalog' import { WsClient } from '@/lib/ws-client' import type { DockerContainer, @@ -20,7 +20,7 @@ const MAX_SECURITY_EVENTS_IN_CACHE = 200 type WsMessage = | { type: 'full_sync'; servers: ServerMetrics[]; upgrades?: UpgradeJob[] } - | { type: 'update'; servers: ServerMetrics[] } + | { type: 'update'; servers: LiveMetrics[] } | { type: 'server_online'; server_id: string } | { type: 'server_offline'; server_id: string } | { diff --git a/apps/web/src/lib/server-catalog.test.ts b/apps/web/src/lib/server-catalog.test.ts index ebdeb3065..a8d6c01a0 100644 --- a/apps/web/src/lib/server-catalog.test.ts +++ b/apps/web/src/lib/server-catalog.test.ts @@ -5,6 +5,7 @@ import { describe, expect, it, vi } from 'vitest' import type { OutstandingEnrollmentSummary, ServerResponse } from '@/lib/api-schema' import { invalidateServerDetail, + type LiveMetrics, projectServerCatalog, readLiveServers, refreshServerCatalog, @@ -138,6 +139,32 @@ function makeMetrics(overrides: Partial = {}): ServerMetrics { } } +function makeLiveMetrics(overrides: Partial = {}): LiveMetrics { + return { + cpu: 42, + disk_read_bytes_per_sec: 100, + disk_used: 125_000, + disk_write_bytes_per_sec: 50, + id: 'server-1', + last_active: 1000, + load1: 1, + load5: 0.8, + load15: 0.6, + mem_used: 16_000, + net_in_speed: 1000, + net_in_transfer: 10_000, + net_out_speed: 500, + net_out_transfer: 5000, + online: true, + process_count: 120, + swap_used: 100, + tcp_conn: 30, + udp_conn: 5, + uptime: 3600, + ...overrides + } +} + const OUTSTANDING_ENROLLMENT: OutstandingEnrollmentSummary = { code_prefix: 'abcdef', created_at: '2026-07-11T00:00:00Z', @@ -192,7 +219,7 @@ describe('server catalog projection', () => { expect(readServerDetail(queryClient, 'server-1')).toEqual(cleared) }) - it('updates runtime metrics without letting WS placeholders erase REST or mutation-owned fields', () => { + it('merges live-only update frames while REST and mutation-owned static fields survive', () => { const queryClient = createQueryClient() const rest = makeRestServer({ has_token: false, @@ -204,26 +231,7 @@ describe('server catalog projection', () => { projectServerCatalog(queryClient, { kind: 'tags_changed', serverId: rest.id, tags: ['prod', 'edge'] }) projectServerCatalog(queryClient, { kind: 'ws_update', - servers: [ - makeMetrics({ - country_code: null, - cpu: 91, - cpu_cores: null, - cpu_name: null, - disk_total: 0, - features: [], - group_id: null, - has_token: true, - mem_total: 0, - mem_used: 24_000, - name: 'Stale Agent Name', - os: null, - outstanding_enrollment: null, - region: null, - swap_total: 0, - tags: [] - }) - ] + servers: [makeLiveMetrics({ cpu: 91, mem_used: 24_000 })] }) expect(readLiveServers(queryClient)?.[0]).toMatchObject({ @@ -246,6 +254,14 @@ describe('server catalog projection', () => { }) }) + it('drops update frames until full sync or REST seeds the catalog', () => { + const queryClient = createQueryClient() + + projectServerCatalog(queryClient, { kind: 'ws_update', servers: [makeLiveMetrics()] }) + + expect(readLiveServers(queryClient)).toBeUndefined() + }) + it('uses full sync as authoritative membership and static state while preserving out-of-band metadata', () => { const queryClient = createQueryClient() const first = makeRestServer({ diff --git a/apps/web/src/lib/server-catalog.ts b/apps/web/src/lib/server-catalog.ts index f43cec607..35d721447 100644 --- a/apps/web/src/lib/server-catalog.ts +++ b/apps/web/src/lib/server-catalog.ts @@ -13,47 +13,56 @@ interface CatalogQueryOptions { enabled?: boolean } -export interface ServerMetrics { - agent_local_capabilities?: number | null - agent_version?: string | null - capabilities?: number - country_code: string | null +/** + * Partial projection carried by WS `update` frames. Mirrors the Rust + * `LiveMetrics` wire type: only fields an agent report can populate. Static + * facts (`mem_total`, `os`, `tags`, ...) never appear here — they arrive via + * `full_sync`/REST and the cached values must survive an update merge. + */ +export interface LiveMetrics { cpu: number - cpu_cores?: number | null - cpu_name: string | null disk_read_bytes_per_sec: number - disk_total: number disk_used: number disk_write_bytes_per_sec: number - effective_capabilities?: number | null - features?: string[] - group_id: string | null - has_token?: boolean id: string last_active: number load1: number load5: number load15: number - mem_total: number mem_used: number - name: string net_in_speed: number net_in_transfer: number net_out_speed: number net_out_transfer: number online: boolean + process_count: number + swap_used: number + tcp_conn: number + udp_conn: number + uptime: number +} + +export interface ServerMetrics extends LiveMetrics { + agent_local_capabilities?: number | null + agent_version?: string | null + capabilities?: number + country_code: string | null + cpu_cores?: number | null + cpu_name: string | null + disk_total: number + effective_capabilities?: number | null + features?: string[] + group_id: string | null + has_token?: boolean + mem_total: number + name: string os: string | null outstanding_enrollment?: OutstandingEnrollmentSummary | null - process_count: number protocol_version?: number region: string | null swap_total: number - swap_used: number tags?: string[] - tcp_conn: number temporary?: TemporaryGrant[] - udp_conn: number - uptime: number } export type ServerCatalogEvent = @@ -68,7 +77,7 @@ export type ServerCatalogEvent = } | { kind: 'tags_changed'; serverId: string; tags: string[] } | { kind: 'ws_full_sync'; servers: ServerMetrics[] } - | { kind: 'ws_update'; servers: ServerMetrics[] } + | { kind: 'ws_update'; servers: LiveMetrics[] } | { kind: 'online_changed'; online: boolean; serverId: string } | { agentLocalCapabilities: number | null | undefined @@ -185,32 +194,11 @@ function projectFullSyncServerToRest(current: ServerResponse, server: ServerMetr } } -function mergeWsServerUpdate(current: ServerMetrics, incoming: ServerMetrics): ServerMetrics { - return { - ...current, - cpu: incoming.cpu, - disk_read_bytes_per_sec: incoming.disk_read_bytes_per_sec, - disk_used: incoming.disk_used, - disk_write_bytes_per_sec: incoming.disk_write_bytes_per_sec, - last_active: incoming.last_active, - load1: incoming.load1, - load5: incoming.load5, - load15: incoming.load15, - mem_used: incoming.mem_used, - net_in_speed: incoming.net_in_speed, - net_in_transfer: incoming.net_in_transfer, - net_out_speed: incoming.net_out_speed, - net_out_transfer: incoming.net_out_transfer, - online: incoming.online, - process_count: incoming.process_count, - swap_used: incoming.swap_used, - tcp_conn: incoming.tcp_conn, - udp_conn: incoming.udp_conn, - uptime: incoming.uptime - } +function mergeWsServerUpdate(current: ServerMetrics, incoming: LiveMetrics): ServerMetrics { + return { ...current, ...incoming } } -function mergeWsUpdate(current: ServerMetrics[], incoming: ServerMetrics[]): ServerMetrics[] { +function mergeWsUpdate(current: ServerMetrics[], incoming: LiveMetrics[]): ServerMetrics[] { const incomingById = new Map(incoming.map((server) => [server.id, server])) return current.map((server) => { const update = incomingById.get(server.id) @@ -427,8 +415,10 @@ export function projectServerCatalog(queryClient: QueryClient, event: ServerCata return } case 'ws_update': { + // An update is a partial projection and can never seed the catalog: + // until full_sync/REST provides the static fields, drop it. queryClient.setQueryData(LIVE_SERVERS_KEY, (current) => - current ? mergeWsUpdate(current, event.servers) : [...event.servers] + current ? mergeWsUpdate(current, event.servers) : current ) return } diff --git a/crates/common/src/protocol.rs b/crates/common/src/protocol.rs index ce20193cd..011905dd6 100644 --- a/crates/common/src/protocol.rs +++ b/crates/common/src/protocol.rs @@ -587,7 +587,7 @@ pub enum BrowserMessage { upgrades: Vec, }, Update { - servers: Vec, + servers: Vec, }, ServerOnline { server_id: String, diff --git a/crates/common/src/types.rs b/crates/common/src/types.rs index 13f66cfac..b53d81226 100644 --- a/crates/common/src/types.rs +++ b/crates/common/src/types.rs @@ -236,6 +236,40 @@ fn default_has_token() -> bool { true } +/// Partial projection of a server carried by `BrowserMessage::Update`. +/// +/// Contains only the fields an agent report can actually populate. Static +/// facts (`mem_total`, `os`, `tags`, ...) are unknown at report time — they +/// travel in the full `ServerStatus` via `FullSync` and REST, and clients keep +/// their cached values when merging an update. Field names serialize +/// identically to `ServerStatus`, so older clients decode an update as a +/// status whose static keys are simply absent. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct LiveMetrics { + pub id: String, + pub online: bool, + pub last_active: i64, + pub uptime: u64, + pub cpu: f64, + pub mem_used: i64, + pub swap_used: i64, + pub disk_used: i64, + pub net_in_speed: i64, + pub net_out_speed: i64, + pub net_in_transfer: i64, + pub net_out_transfer: i64, + pub load1: f64, + pub load5: f64, + pub load15: f64, + pub tcp_conn: i32, + pub udp_conn: i32, + pub process_count: i32, + #[serde(default)] + pub disk_read_bytes_per_sec: u64, + #[serde(default)] + pub disk_write_bytes_per_sec: u64, +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))] pub enum FileType { diff --git a/crates/server/src/service/agent_manager.rs b/crates/server/src/service/agent_manager.rs index d0a5159b7..855328804 100644 --- a/crates/server/src/service/agent_manager.rs +++ b/crates/server/src/service/agent_manager.rs @@ -9,7 +9,7 @@ use tokio::sync::{Mutex, broadcast, mpsc, oneshot}; use serverbee_common::constants::{CAP_DOCKER, has_capability}; use serverbee_common::docker_types::*; use serverbee_common::protocol::{AgentMessage, BrowserMessage, RecordedProtocol, ServerMessage, TemporaryGrant}; -use serverbee_common::types::{ServerStatus, SystemReport, TracerouteHop}; +use serverbee_common::types::{LiveMetrics, SystemReport, TracerouteHop}; use crate::error::AppError; use crate::state::AppState; @@ -247,26 +247,15 @@ impl AgentManager { let (disk_read_bytes_per_sec, disk_write_bytes_per_sec) = aggregate_disk_io(&report); - // Build a ServerStatus for the broadcast. Static fields (mem_total, disk_total, - // os, cpu_name, etc.) are not available here -- set them to defaults since the - // browser can merge with REST data. - let status = ServerStatus { + let metrics = LiveMetrics { id: server_id.to_string(), - name: self - .connections - .get(server_id) - .map(|c| c.server_name.clone()) - .unwrap_or_default(), online: true, last_active: chrono::Utc::now().timestamp(), uptime: report.uptime, cpu: report.cpu, mem_used: report.mem_used, - mem_total: 0, swap_used: report.swap_used, - swap_total: 0, disk_used: report.disk_used, - disk_total: 0, net_in_speed: report.net_in_speed, net_out_speed: report.net_out_speed, net_in_transfer: report.net_in_transfer, @@ -277,22 +266,12 @@ impl AgentManager { tcp_conn: report.tcp_conn, udp_conn: report.udp_conn, process_count: report.process_count, - cpu_name: None, - os: None, - region: None, - country_code: None, - group_id: None, - features: vec![], disk_read_bytes_per_sec, disk_write_bytes_per_sec, - tags: Vec::new(), - cpu_cores: None, - has_token: true, - outstanding_enrollment: None, }; let _ = self.browser_tx.send(BrowserMessage::Update { - servers: vec![status], + servers: vec![metrics], }); // Cache the report @@ -987,6 +966,55 @@ mod tests { assert_eq!(cached.mem_used, 8_000_000_000); } + /// Update frames are a partial projection: static facts must be absent + /// from the wire entirely, so clients keep their cached values on merge + /// instead of having them stomped by zero placeholders. + #[test] + fn test_update_report_broadcast_omits_static_fields() { + let (mgr, mut rx) = make_manager(); + let (tx, _) = mpsc::channel(1); + mgr.add_connection("s1".into(), "Srv".into(), tx, test_addr()); + let _ = rx.try_recv(); // ServerOnline + + mgr.update_report( + "s1", + SystemReport { + cpu: 42.5, + swap_used: 1024, + ..Default::default() + }, + ); + + let msg = rx.try_recv().unwrap(); + let json = serde_json::to_value(&msg).unwrap(); + assert_eq!(json["type"], "update"); + let server = &json["servers"][0]; + assert_eq!(server["id"], "s1"); + assert!((server["cpu"].as_f64().unwrap() - 42.5).abs() < f64::EPSILON); + assert_eq!(server["swap_used"], 1024); + for key in [ + "name", + "mem_total", + "swap_total", + "disk_total", + "cpu_name", + "os", + "region", + "country_code", + "group_id", + "features", + "tags", + "cpu_cores", + "has_token", + "outstanding_enrollment", + ] { + assert!( + server.get(key).is_none(), + "update frame must not carry static field `{key}`" + ); + } + } + #[test] fn test_all_latest_reports() { let (mgr, _rx) = make_manager(); diff --git a/docs/adr/0002-update-frames-carry-live-metrics.md b/docs/adr/0002-update-frames-carry-live-metrics.md new file mode 100644 index 000000000..e06912001 --- /dev/null +++ b/docs/adr/0002-update-frames-carry-live-metrics.md @@ -0,0 +1,26 @@ +# Browser update frames carry LiveMetrics, not a zero-filled ServerStatus + +`BrowserMessage::Update` used to carry the full `ServerStatus` with every static +field (mem_total, os, tags, has_token, ...) hard-coded to `0`/`None`/`true`, +because a `SystemReport` genuinely does not contain them — totals and identity +travel once via `SystemInfo` and live in the DB. The "Update is partial" rule +was enforced nowhere in the types, so every client re-encoded it defensively: +the web merge hand-picked 20 live fields, and iOS guarded `memoryTotal`/`diskTotal` +with `> 0` checks but missed `swapTotal`, which every update frame stomped to 0. + +Decision: update frames carry a dedicated `LiveMetrics` type (Rust +`crates/common/src/types.rs`, mirrored in TS `server-catalog.ts`) containing only +the fields an agent report can populate. `FullSync` and REST keep the full +`ServerStatus`. The wire stays compatible by subtraction — same `update` tag, +same `servers` key, same field names, just fewer keys — so older iOS builds +decode the missing statics as `nil` and their merge keeps cached values (which +also fixes the swapTotal stomp without an app release). + +Consequences: an update can no longer seed the web catalog cache; frames +arriving before `full_sync`/REST are dropped (previously they seeded rows with +zero placeholders — the same bug class client-side). `name` was dropped from +update frames outright: no client consumed it, renames flow through REST. +Considered and rejected: (a) making the static fields `Option` on `ServerStatus` +— one type stays two-faced and every FullSync consumer pays an unwrap; (b) +server-side named constructors only — concentrates the projection but keeps the +zero placeholders on the wire and all the defensive client code alive. From bf6b62ce2f97593fdc0d3ac08ec1815d9b1842a8 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 17:53:47 +0800 Subject: [PATCH 19/44] fix(common): keep name in update frames for shipped iOS decoders Shipped iOS builds decode id and name with non-optional decode() and drop the whole frame when either key is missing, so removing name from LiveMetrics silently killed every live update on existing installs. Keep name as a documented decoder-compat field: it is the agent connection name, not live data. The web merge pins the REST-managed name so the wire value can never stomp an edited one, the server test asserts name is present while static keys stay absent, and a new iOS decoding regression pins the id+name-only requirement against the real LiveMetrics payload. ADR-0002 records the compat floor. --- CONTEXT.md | 2 +- .../BrowserMessageDecodingTests.swift | 55 +++++++++++++++++++ apps/web/src/lib/server-catalog.test.ts | 3 +- apps/web/src/lib/server-catalog.ts | 7 ++- crates/common/src/types.rs | 4 ++ crates/server/src/service/agent_manager.rs | 9 ++- .../0002-update-frames-carry-live-metrics.md | 11 +++- 7 files changed, 84 insertions(+), 7 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index 414cbcf63..4f8b5ec05 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -4,5 +4,5 @@ - **Network tab** — the server-detail tab hosting the network quality detail. Admin gets the full experience (chart, traceroute, target management, CSV export); the public status variant is the redacted summary (targets + anomalies only). - **Server detail tabs** — the top-level structure of the server detail page: Metrics (default, includes the cost/traffic/uptime overview blocks), Network, Traffic, Security, IP Quality. Tab and time range are URL-driven (`?tab=`, `?range=`); the `range` window is shared across tabs. - **Standalone public network page** — `/status/network/$serverId`; exists only as a fallback for status pages configured with `show_server_detail=false` but `show_network=true` (see ADR-0001). -- **Live metrics (实时指标帧)** — the partial projection of a server carried by WS `update` frames: only the fields an agent report can populate (usage, speeds, loads, connection counts). Static facts are absent from the wire; clients keep their cached values when merging (see ADR-0002). _Avoid_: partial status, update payload. +- **Live metrics (实时指标帧)** — the partial projection of a server carried by WS `update` frames: the fields an agent report can populate (usage, speeds, loads, connection counts), plus `name` kept purely for decoder compatibility. Static facts are absent from the wire; clients keep their cached values when merging (see ADR-0002). _Avoid_: partial status, update payload. - **Full server status** — the complete ~35-field snapshot (`ServerStatus`) carried by `full_sync` and REST. The only sources allowed to seed the catalog or overwrite static facts (totals, os, tags, geo, enrollment). diff --git a/apps/ios/ServerBeeTests/BrowserMessageDecodingTests.swift b/apps/ios/ServerBeeTests/BrowserMessageDecodingTests.swift index e7d89df8a..d59734383 100644 --- a/apps/ios/ServerBeeTests/BrowserMessageDecodingTests.swift +++ b/apps/ios/ServerBeeTests/BrowserMessageDecodingTests.swift @@ -95,6 +95,61 @@ final class BrowserMessageDecodingTests: XCTestCase { XCTAssertEqual(servers[0].country, "US") } + /// Regression: since ADR-0002 the server's `update` frame carries the + /// LiveMetrics partial projection — `id` + `name` plus live fields only, + /// with every static key (mem_total, os, tags, has_token, ...) absent. + /// `id` and `name` are the only keys this decoder requires; if either + /// requirement grows, shipped builds silently drop every update frame. + func test_decode_update_acceptsLiveMetricsPartialProjection() throws { + let json = """ + { + "type": "update", + "servers": [ + { + "id": "s1", + "name": "srv-agent-name", + "online": true, + "last_active": 1779834851, + "uptime": 123456, + "cpu": 42.5, + "mem_used": 4294967296, + "swap_used": 1024, + "disk_used": 10737418240, + "net_in_speed": 12345, + "net_out_speed": 67890, + "net_in_transfer": 111, + "net_out_transfer": 222, + "load1": 1.25, + "load5": 0.8, + "load15": 0.6, + "tcp_conn": 34, + "udp_conn": 5, + "process_count": 128, + "disk_read_bytes_per_sec": 100, + "disk_write_bytes_per_sec": 50 + } + ] + } + """ + let msg = try decode(json) + guard case .update(let servers) = msg else { + return XCTFail("Expected .update, got \(msg)") + } + + XCTAssertEqual(servers[0].id, "s1") + XCTAssertEqual(servers[0].cpuUsage, 42.5) + XCTAssertEqual(servers[0].swapUsed, 1024) + XCTAssertEqual(servers[0].uptime, 123_456) + // Static facts are absent from the wire: they must decode to nil so + // merge(from:) keeps the values already learned from full_sync/REST. + XCTAssertNil(servers[0].memoryTotal) + XCTAssertNil(servers[0].swapTotal) + XCTAssertNil(servers[0].diskTotal) + XCTAssertNil(servers[0].os) + XCTAssertNil(servers[0].tags) + XCTAssertNil(servers[0].hasToken) + } + /// Regression: the live browser WS frame sends `last_active` as a Unix epoch /// **integer** (and includes swap/transfer/disk-io/tags/has_token). The old /// decoder typed `last_active` as String, which threw `typeMismatch` and diff --git a/apps/web/src/lib/server-catalog.test.ts b/apps/web/src/lib/server-catalog.test.ts index a8d6c01a0..d453b1f9f 100644 --- a/apps/web/src/lib/server-catalog.test.ts +++ b/apps/web/src/lib/server-catalog.test.ts @@ -151,6 +151,7 @@ function makeLiveMetrics(overrides: Partial = {}): LiveMetrics { load5: 0.8, load15: 0.6, mem_used: 16_000, + name: 'Agent Connection Name', net_in_speed: 1000, net_in_transfer: 10_000, net_out_speed: 500, @@ -231,7 +232,7 @@ describe('server catalog projection', () => { projectServerCatalog(queryClient, { kind: 'tags_changed', serverId: rest.id, tags: ['prod', 'edge'] }) projectServerCatalog(queryClient, { kind: 'ws_update', - servers: [makeLiveMetrics({ cpu: 91, mem_used: 24_000 })] + servers: [makeLiveMetrics({ cpu: 91, mem_used: 24_000, name: 'Stale Agent Name' })] }) expect(readLiveServers(queryClient)?.[0]).toMatchObject({ diff --git a/apps/web/src/lib/server-catalog.ts b/apps/web/src/lib/server-catalog.ts index 35d721447..ef10c85d9 100644 --- a/apps/web/src/lib/server-catalog.ts +++ b/apps/web/src/lib/server-catalog.ts @@ -30,6 +30,8 @@ export interface LiveMetrics { load5: number load15: number mem_used: number + /** Decoder-compat field for shipped iOS builds; never merged (REST owns the name). */ + name: string net_in_speed: number net_in_transfer: number net_out_speed: number @@ -55,7 +57,6 @@ export interface ServerMetrics extends LiveMetrics { group_id: string | null has_token?: boolean mem_total: number - name: string os: string | null outstanding_enrollment?: OutstandingEnrollmentSummary | null protocol_version?: number @@ -195,7 +196,9 @@ function projectFullSyncServerToRest(current: ServerResponse, server: ServerMetr } function mergeWsServerUpdate(current: ServerMetrics, incoming: LiveMetrics): ServerMetrics { - return { ...current, ...incoming } + // The wire `name` exists only so old iOS decoders accept the frame; it is + // the agent connection name and must not stomp the REST-managed one. + return { ...current, ...incoming, name: current.name } } function mergeWsUpdate(current: ServerMetrics[], incoming: LiveMetrics[]): ServerMetrics[] { diff --git a/crates/common/src/types.rs b/crates/common/src/types.rs index b53d81226..a2c67af7a 100644 --- a/crates/common/src/types.rs +++ b/crates/common/src/types.rs @@ -247,6 +247,10 @@ fn default_has_token() -> bool { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct LiveMetrics { pub id: String, + /// Decoder-compat field, not live data: shipped iOS builds require `id` + /// and `name` to decode a frame and drop the whole message otherwise. + /// Clients must NOT let it overwrite the REST-managed name on merge. + pub name: String, pub online: bool, pub last_active: i64, pub uptime: u64, diff --git a/crates/server/src/service/agent_manager.rs b/crates/server/src/service/agent_manager.rs index 855328804..348762349 100644 --- a/crates/server/src/service/agent_manager.rs +++ b/crates/server/src/service/agent_manager.rs @@ -249,6 +249,11 @@ impl AgentManager { let metrics = LiveMetrics { id: server_id.to_string(), + name: self + .connections + .get(server_id) + .map(|c| c.server_name.clone()) + .unwrap_or_default(), online: true, last_active: chrono::Utc::now().timestamp(), uptime: report.uptime, @@ -990,10 +995,12 @@ mod tests { assert_eq!(json["type"], "update"); let server = &json["servers"][0]; assert_eq!(server["id"], "s1"); + // `id` and `name` are the only keys shipped iOS decoders require; a + // frame missing either is dropped wholesale on old builds. + assert_eq!(server["name"], "Srv"); assert!((server["cpu"].as_f64().unwrap() - 42.5).abs() < f64::EPSILON); assert_eq!(server["swap_used"], 1024); for key in [ - "name", "mem_total", "swap_total", "disk_total", diff --git a/docs/adr/0002-update-frames-carry-live-metrics.md b/docs/adr/0002-update-frames-carry-live-metrics.md index e06912001..8bbdb86eb 100644 --- a/docs/adr/0002-update-frames-carry-live-metrics.md +++ b/docs/adr/0002-update-frames-carry-live-metrics.md @@ -16,10 +16,17 @@ same `servers` key, same field names, just fewer keys — so older iOS builds decode the missing statics as `nil` and their merge keeps cached values (which also fixes the swapTotal stomp without an app release). +Subtraction has one hard floor: shipped iOS decoders require `id` **and** +`name` (`ServerStatus.init(from:)` uses non-optional `decode` for both) and +drop the whole frame on a missing key. `LiveMetrics` therefore keeps `name` as +a documented decoder-compat field — it is the agent connection name, not live +data, and clients must not let it overwrite the REST-managed name on merge +(the web merge pins `current.name`; an iOS decoding regression test pins the +contract). + Consequences: an update can no longer seed the web catalog cache; frames arriving before `full_sync`/REST are dropped (previously they seeded rows with -zero placeholders — the same bug class client-side). `name` was dropped from -update frames outright: no client consumed it, renames flow through REST. +zero placeholders — the same bug class client-side). Considered and rejected: (a) making the static fields `Option` on `ServerStatus` — one type stays two-faced and every FullSync consumer pays an unwrap; (b) server-side named constructors only — concentrates the projection but keeps the From d3a3d27c15849736b3b6f855dba8a215ea5ab4bf Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 22:03:23 +0800 Subject: [PATCH 20/44] refactor(server): single-source the streaming upload pending-request keys The upload-ack-{id} / upload-complete-{id} key format was hand-formed in seven places across the HTTP producer (api/file.rs) and the WS consumer (ws/agent/file_transfer.rs); a rename on either side would silently orphan every in-flight upload. The format now lives in one place on AgentManager, next to the pending-request mechanism it keys into. --- crates/server/src/router/api/file.rs | 7 ++++--- .../server/src/router/ws/agent/file_transfer.rs | 9 +++++---- crates/server/src/service/agent_manager.rs | 15 +++++++++++++++ 3 files changed, 24 insertions(+), 7 deletions(-) diff --git a/crates/server/src/router/api/file.rs b/crates/server/src/router/api/file.rs index 98fbbddcb..281ed459f 100644 --- a/crates/server/src/router/api/file.rs +++ b/crates/server/src/router/api/file.rs @@ -14,6 +14,7 @@ use serde::{Deserialize, Serialize}; use crate::error::{ApiResponse, AppError, ok}; use crate::middleware::auth::CurrentUser; +use crate::service::agent_manager::AgentManager; use crate::service::audit::AuditService; use crate::service::capability_gate::require_capability_online; use crate::service::file_transfer::{TransferDirection, TransferInfo}; @@ -911,7 +912,7 @@ async fn upload_file( })?; // Register pending request for the initial ack - let init_ack_key = format!("upload-ack-{transfer_id}"); + let init_ack_key = AgentManager::upload_ack_key(&transfer_id); let init_ack_rx = state .agent_manager .register_pending_request(init_ack_key.clone()); @@ -988,7 +989,7 @@ async fn upload_file( let encoded = base64::engine::general_purpose::STANDARD.encode(&buf[..n]); // Register pending request for the ack (one at a time since upload is sequential) - let ack_msg_id = format!("upload-ack-{transfer_id}"); + let ack_msg_id = AgentManager::upload_ack_key(&transfer_id); let ack_rx = state .agent_manager .register_pending_request(ack_msg_id.clone()); @@ -1059,7 +1060,7 @@ async fn upload_file( .map_err(|_| AppError::Internal("Failed to send upload end".into()))?; // Wait for upload complete or error - let complete_msg_id = format!("upload-complete-{transfer_id}"); + let complete_msg_id = AgentManager::upload_complete_key(&transfer_id); let complete_rx = state .agent_manager .register_pending_request(complete_msg_id); diff --git a/crates/server/src/router/ws/agent/file_transfer.rs b/crates/server/src/router/ws/agent/file_transfer.rs index cf84333c6..328341a6c 100644 --- a/crates/server/src/router/ws/agent/file_transfer.rs +++ b/crates/server/src/router/ws/agent/file_transfer.rs @@ -6,6 +6,7 @@ use std::sync::Arc; +use crate::service::agent_manager::AgentManager; use crate::state::AppState; use serverbee_common::protocol::AgentMessage; @@ -100,7 +101,7 @@ pub(super) fn on_upload_ack( msg: &AgentMessage, ) { state.file_transfers.update_progress(transfer_id, offset); - let ack_key = format!("upload-ack-{transfer_id}"); + let ack_key = AgentManager::upload_ack_key(transfer_id); state .agent_manager .dispatch_pending_response(&ack_key, msg.clone()); @@ -108,7 +109,7 @@ pub(super) fn on_upload_ack( pub(super) fn on_upload_complete(state: &Arc, transfer_id: &str, msg: &AgentMessage) { state.file_transfers.mark_ready(transfer_id); - let complete_key = format!("upload-complete-{transfer_id}"); + let complete_key = AgentManager::upload_complete_key(transfer_id); state .agent_manager .dispatch_pending_response(&complete_key, msg.clone()); @@ -124,8 +125,8 @@ pub(super) fn on_upload_error( .file_transfers .mark_failed(transfer_id, error.to_string()); // The HTTP handler may be waiting on either an ack or complete key — try both. - let ack_key = format!("upload-ack-{transfer_id}"); - let complete_key = format!("upload-complete-{transfer_id}"); + let ack_key = AgentManager::upload_ack_key(transfer_id); + let complete_key = AgentManager::upload_complete_key(transfer_id); if !state .agent_manager .dispatch_pending_response(&complete_key, msg.clone()) diff --git a/crates/server/src/service/agent_manager.rs b/crates/server/src/service/agent_manager.rs index 348762349..68d6087c7 100644 --- a/crates/server/src/service/agent_manager.rs +++ b/crates/server/src/service/agent_manager.rs @@ -457,6 +457,21 @@ impl AgentManager { let _ = self.browser_tx.send(msg); } + /// Pending-request key for a streaming upload's ack exchanges. Uploads + /// cannot use the one-shot `request()` seam (init-ack, per-chunk ack and + /// complete are separate exchanges over one transfer id), so the HTTP + /// producer and the WS consumer must form identical keys — the format + /// lives here and nowhere else. + pub fn upload_ack_key(transfer_id: &str) -> String { + format!("upload-ack-{transfer_id}") + } + + /// Pending-request key for a streaming upload's completion. See + /// [`Self::upload_ack_key`]. + pub fn upload_complete_key(transfer_id: &str) -> String { + format!("upload-complete-{transfer_id}") + } + /// Register a pending request for HTTP→WS relay with a custom TTL. /// Returns a oneshot receiver that will receive the agent's response. pub fn register_pending_request_with_ttl( From e3e520311e79a946087db97a7a4562550344a424 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 22:03:23 +0800 Subject: [PATCH 21/44] refactor(server): unify the two IP-change reaction paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit on_system_info and on_ip_changed each hand-rolled the same audit + alert + ServerIpChanged broadcast sequence with slightly different inputs, so a fix to one path could silently miss the other. Both now funnel through apply_ip_change with an explicit IpChange struct; persistence stays with each caller. Behavior deltas: agent-IP (v4/v6) transitions detected at SystemInfo now write an audit row too (previously only remote-addr changes did), and first population (None -> Some) never audits — it happens on every fresh registration and only fed noise. Alert and broadcast semantics are unchanged, including the tolerated first-connect frames. --- .../server/src/router/ws/agent/system_info.rs | 217 ++++++++++-------- 1 file changed, 125 insertions(+), 92 deletions(-) diff --git a/crates/server/src/router/ws/agent/system_info.rs b/crates/server/src/router/ws/agent/system_info.rs index 299ac9132..96e33ae25 100644 --- a/crates/server/src/router/ws/agent/system_info.rs +++ b/crates/server/src/router/ws/agent/system_info.rs @@ -38,6 +38,111 @@ fn resolve_public_ip( .find(|ip| !ip.is_loopback() && !geoip::is_private(ip)) } +/// A detected change of a server's externally visible addresses. +struct IpChange { + old_ipv4: Option, + new_ipv4: Option, + old_ipv6: Option, + new_ipv6: Option, + old_remote_addr: Option, + new_remote_addr: Option, +} + +impl IpChange { + /// Any difference at all, including first population (`None` → `Some`). + /// Drives the alert check and the browser broadcast. + fn changed(&self) -> bool { + self.old_ipv4 != self.new_ipv4 + || self.old_ipv6 != self.new_ipv6 + || self.old_remote_addr != self.new_remote_addr + } + + /// A real transition (`Some` → different value) on at least one field. + /// Drives the audit trail: first population happens on every fresh + /// registration and must not spam the audit log. + fn is_transition(&self) -> bool { + fn t(old: &Option, new: &Option) -> bool { + old.is_some() && old != new + } + t(&self.old_ipv4, &self.new_ipv4) + || t(&self.old_ipv6, &self.new_ipv6) + || t(&self.old_remote_addr, &self.new_remote_addr) + } + + fn detail(&self, server_id: &str) -> String { + let mut parts = Vec::new(); + if self.old_ipv4 != self.new_ipv4 { + parts.push(format!("ipv4 {:?} -> {:?}", self.old_ipv4, self.new_ipv4)); + } + if self.old_ipv6 != self.new_ipv6 { + parts.push(format!("ipv6 {:?} -> {:?}", self.old_ipv6, self.new_ipv6)); + } + if self.old_remote_addr != self.new_remote_addr { + parts.push(format!( + "remote {:?} -> {:?}", + self.old_remote_addr, self.new_remote_addr + )); + } + format!("IP changed for server {server_id}: {}", parts.join(", ")) + } +} + +/// The single reaction to a detected IP change, shared by the `SystemInfo` +/// and `IpChanged` paths: audit trail, alert event rules, browser broadcast. +/// Persistence stays with the caller — `SystemInfo` persists the whole +/// identity row via `update_system_info`, `IpChanged` updates just the IP +/// columns — so this function must never write address columns itself. +async fn apply_ip_change(state: &Arc, server_id: &str, change: IpChange) { + if !change.changed() { + return; + } + + if change.is_transition() { + let detail = change.detail(server_id); + tracing::info!("{detail}"); + + let audit_ip = change + .new_remote_addr + .clone() + .or_else(|| { + state + .agent_manager + .get_remote_addr(server_id) + .map(|a| a.ip().to_string()) + }) + .unwrap_or_default(); + if let Err(e) = + AuditService::log(&state.db, "system", "ip_changed", Some(&detail), &audit_ip).await + { + tracing::error!("Failed to write audit log for IP change: {e}"); + } + } + + if let Err(e) = AlertService::check_event_rules( + &state.db, + &state.config, + &state.alert_state_manager, + server_id, + "ip_changed", + ) + .await + { + tracing::error!("Failed to check event rules for IP change: {e}"); + } + + state + .agent_manager + .broadcast_browser(BrowserMessage::ServerIpChanged { + server_id: server_id.to_string(), + old_ipv4: change.old_ipv4, + new_ipv4: change.new_ipv4, + old_ipv6: change.old_ipv6, + new_ipv6: change.new_ipv6, + old_remote_addr: change.old_remote_addr, + new_remote_addr: change.new_remote_addr, + }); +} + pub(super) async fn on_system_info( state: &Arc, server_id: &str, @@ -77,61 +182,19 @@ pub(super) async fn on_system_info( .map(|a| a.ip().to_string()); if let Ok(srv) = ServerService::get_server(&state.db, server_id).await { - let old_remote_addr = srv.last_remote_addr.clone(); - let old_ipv4 = srv.ipv4.clone(); - let old_ipv6 = srv.ipv6.clone(); - - // Check if remote_addr changed - if let Some(ref new_addr) = current_remote_addr - && let Some(ref old_addr) = old_remote_addr - && old_addr != new_addr - { - tracing::info!("Server {server_id} remote address changed: {old_addr} -> {new_addr}"); - if let Err(e) = AuditService::log( - &state.db, - "system", - "ip_changed", - Some(&format!( - "Remote address changed from {old_addr} to {new_addr} for server {server_id}" - )), - new_addr, - ) - .await - { - tracing::error!("Failed to write audit log for IP change: {e}"); - } - } - - // Check if agent-reported IPs changed - let ipv4_changed = old_ipv4 != info.ipv4; - let ipv6_changed = old_ipv6 != info.ipv6; - let remote_changed = old_remote_addr.as_ref() != current_remote_addr.as_ref(); - - if ipv4_changed || ipv6_changed || remote_changed { - if let Err(e) = AlertService::check_event_rules( - &state.db, - &state.config, - &state.alert_state_manager, - server_id, - "ip_changed", - ) - .await - { - tracing::error!("Failed to check event rules for IP change: {e}"); - } - - state - .agent_manager - .broadcast_browser(BrowserMessage::ServerIpChanged { - server_id: server_id.to_string(), - old_ipv4, - new_ipv4: info.ipv4.clone(), - old_ipv6, - new_ipv6: info.ipv6.clone(), - old_remote_addr, - new_remote_addr: current_remote_addr.clone(), - }); - } + apply_ip_change( + state, + server_id, + IpChange { + old_ipv4: srv.ipv4.clone(), + new_ipv4: info.ipv4.clone(), + old_ipv6: srv.ipv6.clone(), + new_ipv6: info.ipv6.clone(), + old_remote_addr: srv.last_remote_addr.clone(), + new_remote_addr: current_remote_addr.clone(), + }, + ) + .await; // Always update last_remote_addr if let Some(ref addr) = current_remote_addr @@ -269,10 +332,8 @@ pub(super) async fn on_ip_changed( Ok(srv) => { let old_ipv4 = srv.ipv4.clone(); let old_ipv6 = srv.ipv6.clone(); - let ipv4_changed = old_ipv4 != ipv4; - let ipv6_changed = old_ipv6 != ipv6; - if ipv4_changed || ipv6_changed { + if old_ipv4 != ipv4 || old_ipv6 != ipv6 { // Update ipv4/ipv6 in DB if let Err(e) = update_server_ips(&state.db, server_id, &ipv4, &ipv6).await { tracing::error!("Failed to update IPs for {server_id}: {e}"); @@ -297,47 +358,19 @@ pub(super) async fn on_ip_changed( } } - let detail = format!( - "IP changed for server {server_id}: ipv4 {:?} -> {:?}, ipv6 {:?} -> {:?}", - old_ipv4, ipv4, old_ipv6, ipv6 - ); - tracing::info!("{detail}"); - - let remote_ip = state - .agent_manager - .get_remote_addr(server_id) - .map(|a| a.ip().to_string()) - .unwrap_or_default(); - if let Err(e) = - AuditService::log(&state.db, "system", "ip_changed", Some(&detail), &remote_ip) - .await - { - tracing::error!("Failed to write audit log for IP change: {e}"); - } - - if let Err(e) = AlertService::check_event_rules( - &state.db, - &state.config, - &state.alert_state_manager, + apply_ip_change( + state, server_id, - "ip_changed", - ) - .await - { - tracing::error!("Failed to check event rules for IP change: {e}"); - } - - state - .agent_manager - .broadcast_browser(BrowserMessage::ServerIpChanged { - server_id: server_id.to_string(), + IpChange { old_ipv4, new_ipv4: ipv4, old_ipv6, new_ipv6: ipv6, old_remote_addr: None, new_remote_addr: None, - }); + }, + ) + .await; } } Err(e) => { From ba27e1482f41ff1263ecc1d57acfc400ef00fc13 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 22:03:23 +0800 Subject: [PATCH 22/44] refactor(web): route network_probe_update through subscribeBrowserMessage The WS layer had two pub/sub seams for the same job: traceroute frames used the module-owned subscribeBrowserMessage registry while network_probe_update detoured through a global window CustomEvent that no test could reach. useNetworkRealtime now subscribes like every other consumer, and the branch gains dispatch + malformed-frame regression tests. --- apps/web/src/hooks/use-network-realtime.ts | 14 ++++----- apps/web/src/hooks/use-servers-ws.test.ts | 35 +++++++++++++++++++++- apps/web/src/hooks/use-servers-ws.ts | 7 +---- 3 files changed, 41 insertions(+), 15 deletions(-) diff --git a/apps/web/src/hooks/use-network-realtime.ts b/apps/web/src/hooks/use-network-realtime.ts index bedc7e3d2..36149b8f2 100644 --- a/apps/web/src/hooks/use-network-realtime.ts +++ b/apps/web/src/hooks/use-network-realtime.ts @@ -1,4 +1,5 @@ import { useCallback, useEffect, useRef, useState } from 'react' +import { subscribeBrowserMessage } from '@/hooks/use-servers-ws' import type { NetworkProbeResultData } from '@/lib/network-types' const MAX_POINTS = 200 @@ -13,13 +14,13 @@ export function useNetworkRealtime(serverId: string) { const dataRef = useRef({}) const handleUpdate = useCallback( - (event: Event) => { - const detail = (event as CustomEvent).detail - if (detail.server_id !== serverId) { + (msg: Record) => { + if (msg.server_id !== serverId) { return } - const results: NetworkProbeResultData[] = detail.results + // handleWsMessage validated the shape before dispatching. + const results = msg.results as NetworkProbeResultData[] const newData = { ...dataRef.current } for (const result of results) { @@ -36,10 +37,7 @@ export function useNetworkRealtime(serverId: string) { [serverId] ) - useEffect(() => { - window.addEventListener('network-probe-update', handleUpdate) - return () => window.removeEventListener('network-probe-update', handleUpdate) - }, [handleUpdate]) + useEffect(() => subscribeBrowserMessage('network_probe_update', handleUpdate), [handleUpdate]) const reset = useCallback(() => { dataRef.current = {} diff --git a/apps/web/src/hooks/use-servers-ws.test.ts b/apps/web/src/hooks/use-servers-ws.test.ts index 665c8e7ae..5513e22ef 100644 --- a/apps/web/src/hooks/use-servers-ws.test.ts +++ b/apps/web/src/hooks/use-servers-ws.test.ts @@ -1,7 +1,7 @@ import { QueryClient } from '@tanstack/react-query' import { describe, expect, it } from 'vitest' import { useUpgradeJobsStore } from '@/stores/upgrade-jobs-store' -import { handleWsMessage } from './use-servers-ws' +import { handleWsMessage, subscribeBrowserMessage } from './use-servers-ws' describe('handleWsMessage upgrade messages', () => { it('hydrates upgrade jobs from full_sync', () => { @@ -92,3 +92,36 @@ describe('handleWsMessage upgrade messages', () => { expect(job?.finished_at).not.toBeNull() }) }) + +describe('handleWsMessage network probe updates', () => { + it('dispatches validated network_probe_update frames to subscribers', () => { + const queryClient = new QueryClient() + const received: Array> = [] + const unsubscribe = subscribeBrowserMessage('network_probe_update', (msg) => received.push(msg)) + + handleWsMessage( + { + type: 'network_probe_update', + server_id: 's1', + results: [{ latency_ms: 12, target_id: 't1' }] + }, + queryClient + ) + + unsubscribe() + expect(received).toHaveLength(1) + expect(received[0].server_id).toBe('s1') + }) + + it('drops malformed network_probe_update frames before dispatch', () => { + const queryClient = new QueryClient() + const received: Array> = [] + const unsubscribe = subscribeBrowserMessage('network_probe_update', (msg) => received.push(msg)) + + handleWsMessage({ results: [null], server_id: 's1', type: 'network_probe_update' }, queryClient) + handleWsMessage({ results: [{}], type: 'network_probe_update' }, queryClient) + + unsubscribe() + expect(received).toHaveLength(0) + }) +}) diff --git a/apps/web/src/hooks/use-servers-ws.ts b/apps/web/src/hooks/use-servers-ws.ts index 868fa0cfb..85f945fd5 100644 --- a/apps/web/src/hooks/use-servers-ws.ts +++ b/apps/web/src/hooks/use-servers-ws.ts @@ -406,12 +406,7 @@ export function handleWsMessage(raw: unknown, queryClient: QueryClient): void { ) { break } - const msg = raw as WsMessage & { type: 'network_probe_update' } - window.dispatchEvent( - new CustomEvent('network-probe-update', { - detail: { server_id: msg.server_id, results: msg.results } - }) - ) + dispatchToSubscribers('network_probe_update', raw) break } case 'docker_update': From 58c47d3763b87acbe5eebcca1fbf41a109db58be Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 22:19:27 +0800 Subject: [PATCH 23/44] refactor(server): share the browser-WS gate and mobile-expiry scaffold terminal.rs, docker_logs.rs and browser.rs each re-implemented the same session scaffold: credential resolve, admin-role denial audit, online check, audited capability gate, and a byte-identical mobile-token-expiry select! arm. A fix to any of those had to be applied three times. ws/session.rs now owns both pieces: admin_capability_gate runs the whole pre-upgrade sequence and returns what the pump needs (user, ip, mobile deadline), and mobile_token_expired is the shared expiry future. Routes keep only their own pumps; the open/close audit bookends stay per-route because their contexts genuinely differ. --- crates/server/src/router/ws/browser.rs | 10 +-- crates/server/src/router/ws/docker_logs.rs | 89 ++++++---------------- crates/server/src/router/ws/mod.rs | 1 + crates/server/src/router/ws/session.rs | 88 +++++++++++++++++++++ crates/server/src/router/ws/terminal.rs | 83 ++++++-------------- 5 files changed, 137 insertions(+), 134 deletions(-) create mode 100644 crates/server/src/router/ws/session.rs diff --git a/crates/server/src/router/ws/browser.rs b/crates/server/src/router/ws/browser.rs index 4e7b7bf46..736e3a933 100644 --- a/crates/server/src/router/ws/browser.rs +++ b/crates/server/src/router/ws/browser.rs @@ -115,15 +115,7 @@ async fn handle_browser_ws( } } } - // Mobile token expiry: auto-close when the token expires - _ = async { - if let Some(exp) = mobile_expires { - let dur = (exp - chrono::Utc::now()).to_std().unwrap_or_default(); - tokio::time::sleep(dur).await; - } else { - std::future::pending::<()>().await; - } - } => { + () = super::session::mobile_token_expired(mobile_expires) => { tracing::debug!("Mobile WS token expired, closing connection"); let _ = ws_sink.send(Message::Close(Some(axum::extract::ws::CloseFrame { code: 4001, diff --git a/crates/server/src/router/ws/docker_logs.rs b/crates/server/src/router/ws/docker_logs.rs index db78c67e8..e37a6b850 100644 --- a/crates/server/src/router/ws/docker_logs.rs +++ b/crates/server/src/router/ws/docker_logs.rs @@ -4,14 +4,12 @@ use axum::Router; use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade}; use axum::extract::{ConnectInfo, Path, State}; use axum::http::HeaderMap; -use axum::response::{IntoResponse, Response}; +use axum::response::Response; use axum::routing::get; use futures_util::{SinkExt, StreamExt}; use serde::Deserialize; use tokio::sync::mpsc; -use crate::middleware::auth::resolve_ws_connection; -use crate::router::utils::extract_client_ip; use crate::service::audit::AuditService; use crate::service::high_risk_audit::DockerLogsAuditContext; use crate::state::AppState; @@ -29,60 +27,32 @@ async fn docker_logs_ws_handler( headers: HeaderMap, ws: WebSocketUpgrade, ) -> Response { - let ip = extract_client_ip( - &ConnectInfo(addr), + // Docker log streaming exposes sensitive container output (env vars, + // connection strings, tokens), so it is admin-only like the terminal; + // the shared gate audits denials under `docker_logs_subscribe_denied`. + match super::session::admin_capability_gate( + &state, &headers, - &state.config.server.trusted_proxies, + &ConnectInfo(addr), + &server_id, + CAP_DOCKER, + "docker_logs_subscribe_denied", ) - .to_string(); - - // Shared credential policy lives in `middleware::auth`; this adapter only - // applies Docker-specific rules (admin role, agent online, capability). - match resolve_ws_connection(&headers, &state).await { - Some(conn) => { - let user_id = conn.user.user_id; - // Docker log streaming exposes sensitive container output - // (env vars, connection strings, tokens), so it is admin-only, - // consistent with terminal access. - if conn.user.role != "admin" { - let detail = serde_json::json!({ - "server_id": server_id, - "deny_reason": "role_forbidden", - }) - .to_string(); - let _ = AuditService::log( - &state.db, - &user_id, - "docker_logs_subscribe_denied", - Some(&detail), - &ip, + .await + { + Ok(gate) => ws + .max_message_size(MAX_WS_MESSAGE_SIZE) + .on_upgrade(move |socket| { + handle_docker_logs_ws( + socket, + state, + server_id, + gate.user_id, + gate.ip, + gate.mobile_expires, ) - .await; - return axum::http::StatusCode::FORBIDDEN.into_response(); - } - // Check agent is online - if !state.agent_manager.is_online(&server_id) { - return (axum::http::StatusCode::BAD_REQUEST, "Agent is offline").into_response(); - } - // Check Docker capability (denials are audited by the gate) - if let Err(error) = crate::service::capability_gate::require_capability_audited( - &state, - &server_id, - CAP_DOCKER, - &user_id, - &ip, - "docker_logs_subscribe_denied", - ) - .await - { - return error.into_response(); - } - ws.max_message_size(MAX_WS_MESSAGE_SIZE) - .on_upgrade(move |socket| { - handle_docker_logs_ws(socket, state, server_id, user_id, ip, conn.mobile_expires) - }) - } - None => axum::http::StatusCode::UNAUTHORIZED.into_response(), + }), + Err(response) => response, } } @@ -227,16 +197,7 @@ async fn handle_docker_logs_ws( } } } - // Mobile token expiry: force-close when a fixed-lifetime mobile token - // expires mid-session (web sessions / API keys never trip this arm). - () = async { - if let Some(exp) = mobile_expires { - let dur = (exp - chrono::Utc::now()).to_std().unwrap_or_default(); - tokio::time::sleep(dur).await; - } else { - std::future::pending::<()>().await; - } - } => { + () = super::session::mobile_token_expired(mobile_expires) => { tracing::debug!("Docker logs session {session_id} mobile token expired, closing"); break "token_expired"; } diff --git a/crates/server/src/router/ws/mod.rs b/crates/server/src/router/ws/mod.rs index 6e62b46d8..8d3a1d9cf 100644 --- a/crates/server/src/router/ws/mod.rs +++ b/crates/server/src/router/ws/mod.rs @@ -1,4 +1,5 @@ pub mod agent; pub mod browser; pub mod docker_logs; +mod session; pub mod terminal; diff --git a/crates/server/src/router/ws/session.rs b/crates/server/src/router/ws/session.rs new file mode 100644 index 000000000..ae1aa3042 --- /dev/null +++ b/crates/server/src/router/ws/session.rs @@ -0,0 +1,88 @@ +//! Shared session scaffold for browser-facing WS routes. +//! +//! Every control-plane WS route must make the same security-sensitive +//! decisions before upgrading (credential policy → admin role → agent online +//! → capability, denials audited) and watch the same mobile-token deadline +//! while pumping. Both live here so a fix applies to every route at once; +//! the routes keep only their own pumps. + +use std::sync::Arc; + +use axum::extract::ConnectInfo; +use axum::http::HeaderMap; +use axum::response::{IntoResponse, Response}; + +use crate::middleware::auth::resolve_ws_connection; +use crate::router::utils::extract_client_ip; +use crate::service::audit::AuditService; +use crate::service::capability_gate::require_capability_audited; +use crate::state::AppState; + +/// What a gated WS pump needs from the pre-upgrade checks. +pub(super) struct WsGate { + pub user_id: String, + pub ip: String, + pub mobile_expires: Option>, +} + +/// Admin gate for control-plane WS routes (terminal, docker logs). +/// +/// Runs the shared credential policy, then requires the admin role, an online +/// agent, and `capability` — auditing role/capability denials under +/// `denied_action`. Returns the ready-to-serve response on any failure so the +/// route handler can simply `match`. +pub(super) async fn admin_capability_gate( + state: &Arc, + headers: &HeaderMap, + addr: &ConnectInfo, + server_id: &str, + capability: u32, + denied_action: &str, +) -> Result { + let ip = extract_client_ip(addr, headers, &state.config.server.trusted_proxies).to_string(); + + let Some(conn) = resolve_ws_connection(headers, state).await else { + return Err(axum::http::StatusCode::UNAUTHORIZED.into_response()); + }; + let user_id = conn.user.user_id; + + if conn.user.role != "admin" { + let detail = serde_json::json!({ + "server_id": server_id, + "deny_reason": "role_forbidden", + }) + .to_string(); + let _ = AuditService::log(&state.db, &user_id, denied_action, Some(&detail), &ip).await; + return Err(axum::http::StatusCode::FORBIDDEN.into_response()); + } + + if !state.agent_manager.is_online(server_id) { + return Err((axum::http::StatusCode::BAD_REQUEST, "Agent is offline").into_response()); + } + + if let Err(error) = + require_capability_audited(state, server_id, capability, &user_id, &ip, denied_action) + .await + { + return Err(error.into_response()); + } + + Ok(WsGate { + user_id, + ip, + mobile_expires: conn.mobile_expires, + }) +} + +/// Resolves when a fixed-lifetime mobile token expires; pends forever for web +/// sessions and API keys. The shared `select!` arm of every browser-facing +/// WS pump. +pub(super) async fn mobile_token_expired(expires: Option>) { + match expires { + Some(exp) => { + let dur = (exp - chrono::Utc::now()).to_std().unwrap_or_default(); + tokio::time::sleep(dur).await; + } + None => std::future::pending::<()>().await, + } +} diff --git a/crates/server/src/router/ws/terminal.rs b/crates/server/src/router/ws/terminal.rs index fdb0da14b..899ae5379 100644 --- a/crates/server/src/router/ws/terminal.rs +++ b/crates/server/src/router/ws/terminal.rs @@ -4,14 +4,12 @@ use axum::Router; use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade}; use axum::extract::{ConnectInfo, Path, State}; use axum::http::HeaderMap; -use axum::response::{IntoResponse, Response}; +use axum::response::Response; use axum::routing::get; use futures_util::{SinkExt, StreamExt}; use serde::Deserialize; use tokio::sync::mpsc; -use crate::middleware::auth::resolve_ws_connection; -use crate::router::utils::extract_client_ip; use crate::service::agent_manager::TerminalSessionEvent; use crate::service::audit::AuditService; use crate::service::high_risk_audit::TerminalAuditContext; @@ -30,57 +28,29 @@ async fn terminal_ws_handler( headers: HeaderMap, ws: WebSocketUpgrade, ) -> Response { - let ip = extract_client_ip( - &ConnectInfo(addr), + // Terminal is admin-only and needs the terminal capability; the shared + // gate audits denials under `terminal_open_denied`. + match super::session::admin_capability_gate( + &state, &headers, - &state.config.server.trusted_proxies, + &ConnectInfo(addr), + &server_id, + serverbee_common::constants::CAP_TERMINAL, + "terminal_open_denied", ) - .to_string(); - - // Shared credential policy lives in `middleware::auth`; this adapter only - // applies terminal-specific rules (admin role, agent online, capability). - match resolve_ws_connection(&headers, &state).await { - Some(conn) => { - let user_id = conn.user.user_id; - // Terminal access is admin-only - if conn.user.role != "admin" { - let detail = serde_json::json!({ - "server_id": server_id, - "deny_reason": "role_forbidden", - }) - .to_string(); - let _ = AuditService::log( - &state.db, - &user_id, - "terminal_open_denied", - Some(&detail), - &ip, - ) - .await; - return axum::http::StatusCode::FORBIDDEN.into_response(); - } - // Check agent is online - if !state.agent_manager.is_online(&server_id) { - return (axum::http::StatusCode::BAD_REQUEST, "Agent is offline").into_response(); - } - // Check terminal capability (denials are audited by the gate) - if let Err(error) = crate::service::capability_gate::require_capability_audited( - &state, - &server_id, - serverbee_common::constants::CAP_TERMINAL, - &user_id, - &ip, - "terminal_open_denied", + .await + { + Ok(gate) => ws.max_message_size(MAX_WS_MESSAGE_SIZE).on_upgrade(move |socket| { + handle_terminal_ws( + socket, + state, + server_id, + gate.user_id, + gate.ip, + gate.mobile_expires, ) - .await - { - return error.into_response(); - } - ws.max_message_size(MAX_WS_MESSAGE_SIZE).on_upgrade(move |socket| { - handle_terminal_ws(socket, state, server_id, user_id, ip, conn.mobile_expires) - }) - } - None => axum::http::StatusCode::UNAUTHORIZED.into_response(), + }), + Err(response) => response, } } @@ -255,16 +225,7 @@ async fn handle_terminal_ws( close_reason = "idle_timeout".to_string(); break; } - // Mobile token expiry: force-close when a fixed-lifetime mobile token - // expires mid-session (web sessions / API keys never trip this arm). - () = async { - if let Some(exp) = mobile_expires { - let dur = (exp - chrono::Utc::now()).to_std().unwrap_or_default(); - tokio::time::sleep(dur).await; - } else { - std::future::pending::<()>().await; - } - } => { + () = super::session::mobile_token_expired(mobile_expires) => { tracing::debug!("Terminal session {session_id} mobile token expired, closing"); let msg = serde_json::json!({"type": "error", "error": "Session token expired"}); let _ = ws_sink.send(Message::Text(msg.to_string().into())).await; From e3f190f0c260c7e3a651dbcd37d7ac7123ad3761 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 22:39:41 +0800 Subject: [PATCH 24/44] refactor(web): extract the Metrics-tab chart model into lib/metric-chart-model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit server-detail-content.tsx buried every chart transform — the metric row mappers, the pct guard, realtime tick dedup, long-range tick and tooltip formats, the x-axis stride and GPU row shaping — as unexported inner functions, so none of it was reachable by a test and the row mapper existed three times (admin, public, realtime). lib/metric-chart-model.ts now owns the transforms behind small pure functions: one toMetricChartRow covers both admin records and public points, one pct guard is shared with toRealtimeDataPoint, and the formatters take plain rangeHours. The component keeps rendering plus thin useMemo glue. 22 new test cases cover the model; chart behavior verified in the browser across realtime/24h/7d ranges. --- .../status/server-detail-content.tsx | 180 ++------------- apps/web/src/hooks/use-realtime-metrics.ts | 5 +- apps/web/src/lib/metric-chart-model.test.ts | 209 ++++++++++++++++++ apps/web/src/lib/metric-chart-model.ts | 208 +++++++++++++++++ 4 files changed, 439 insertions(+), 163 deletions(-) create mode 100644 apps/web/src/lib/metric-chart-model.test.ts create mode 100644 apps/web/src/lib/metric-chart-model.ts diff --git a/apps/web/src/components/status/server-detail-content.tsx b/apps/web/src/components/status/server-detail-content.tsx index 31067458c..6e0273775 100644 --- a/apps/web/src/components/status/server-detail-content.tsx +++ b/apps/web/src/components/status/server-detail-content.tsx @@ -25,19 +25,20 @@ import type { UptimeDailyEntry } from '@/lib/api-schema' import { buildMergedDiskIoSeries, buildPerDiskIoSeries } from '@/lib/disk-io' -import { type ServerMetrics, useLiveServers } from '@/lib/server-catalog' +import { + buildGpuChartRows, + deriveNetworkLabels, + type GpuRecordAggregated, + makeTickFormatter, + makeTooltipFormatter, + toMetricChartRow, + xAxisStride +} from '@/lib/metric-chart-model' +import { useLiveServers } from '@/lib/server-catalog' import { type RangeKey, rangesForVariant, resolveRange, type TimeRange } from '@/lib/server-detail-nav' import { cn, formatBytes } from '@/lib/utils' import { computeAggregateUptime } from '@/lib/widget-helpers' -interface GpuRecordAggregated { - gpu_usage_avg: number - mem_total_avg: number - mem_used_avg: number - temperature_avg: number - time: string -} - export interface ServerDetailContentProps { /** Currently selected detail tab. When provided (admin), the tabs become * URL-controlled; the public surface omits it and stays uncontrolled. */ @@ -73,40 +74,6 @@ function buildIsoWindow(hours: number) { } } -function adminRecordToChartRow(r: ServerMetricRecord, memTotal: number, diskTotal: number) { - return { - timestamp: r.time, - cpu: r.cpu, - memory_pct: memTotal ? (r.mem_used / memTotal) * 100 : 0, - disk_pct: diskTotal ? (r.disk_used / diskTotal) * 100 : 0, - net_in_speed: r.net_in_speed, - net_out_speed: r.net_out_speed, - net_in_transfer: r.net_in_transfer, - net_out_transfer: r.net_out_transfer, - load1: r.load1, - load5: r.load5, - load15: r.load15, - temperature: r.temperature - } -} - -function publicPointToChartRow(p: PublicMetricsPoint, memTotal: number, diskTotal: number) { - return { - timestamp: p.time, - cpu: p.cpu, - memory_pct: memTotal ? (p.mem_used / memTotal) * 100 : 0, - disk_pct: diskTotal ? (p.disk_used / diskTotal) * 100 : 0, - net_in_speed: p.net_in_speed, - net_out_speed: p.net_out_speed, - net_in_transfer: p.net_in_transfer, - net_out_transfer: p.net_out_transfer, - load1: p.load1, - load5: p.load5, - load15: p.load15, - temperature: p.temperature - } -} - // Fetches the historical metric series, branching on variant. Admin uses the // auth'd `useServerRecords` (includes disk-io + temperature blobs); public // hits `/api/status/servers/{id}/metrics` which returns the normalised @@ -151,45 +118,6 @@ function useLiveServerMetrics(serverId: string, isAdminVariant: boolean) { return liveServers?.find((s) => s.id === serverId) } -interface NetworkLabels { - netInLabel: string - netOutLabel: string - netTotalLabel: string | null -} - -function deriveNetworkLabels( - isAdminVariant: boolean, - liveData: ServerMetrics | undefined, - publicMetricsSnapshot: PublicServerDetail['metrics'] | null -): NetworkLabels { - if (isAdminVariant) { - if (!liveData) { - return { netInLabel: '—', netOutLabel: '—', netTotalLabel: '—' } - } - const inBytes = liveData.net_in_transfer ?? 0 - const outBytes = liveData.net_out_transfer ?? 0 - return { - netInLabel: formatBytes(inBytes), - netOutLabel: formatBytes(outBytes), - netTotalLabel: formatBytes(inBytes + outBytes) - } - } - if (!publicMetricsSnapshot) { - return { netInLabel: '—', netOutLabel: '—', netTotalLabel: null } - } - // Use cumulative transfer (not the instantaneous *_speed rate) so the public - // bar matches the admin bar: the `detail_network_in/out/total` labels describe a - // total amount transferred, and formatBytes renders bytes — feeding a rate here - // mislabelled "1.2 MB/s" as a cumulative "1.2 MB". - const inBytes = publicMetricsSnapshot.net_in_transfer - const outBytes = publicMetricsSnapshot.net_out_transfer - return { - netInLabel: formatBytes(inBytes), - netOutLabel: formatBytes(outBytes), - netTotalLabel: formatBytes(inBytes + outBytes) - } -} - export function ServerDetailContent(props: ServerDetailContentProps) { const { activeTab, networkTab, rangeKey, server, serverId, onRangeChange, onTabChange, variant } = props const { t } = useTranslation('servers') @@ -411,79 +339,22 @@ function useAggregatedChartData(args: { if (isRealtime) { return realtimeData as Record[] } - return (adminRecords ?? []).map((r) => adminRecordToChartRow(r, memTotal, diskTotal)) + return (adminRecords ?? []).map((r) => toMetricChartRow(r, memTotal, diskTotal)) } - return (publicMetrics ?? []).map((p) => publicPointToChartRow(p, memTotal, diskTotal)) + return (publicMetrics ?? []).map((p) => toMetricChartRow(p, memTotal, diskTotal)) }, [adminRecords, diskTotal, isAdminVariant, isRealtime, memTotal, publicMetrics, realtimeData]) } -function formatHourMinute(time: string) { - const d = new Date(time) - return `${String(d.getHours()).padStart(2, '0')}:${String(d.getMinutes()).padStart(2, '0')}` -} - function useChartTickFormatter(isRealtime: boolean, range: TimeRange, chartData: Record[]) { - return useMemo<((time: string) => string) | undefined>(() => { - if (isRealtime) { - const realtimeLabels = new Map() - let previousLabel = '' - for (const point of chartData) { - if (typeof point.timestamp !== 'string') { - continue - } - const label = formatHourMinute(point.timestamp) - realtimeLabels.set(point.timestamp, label === previousLabel ? '' : label) - previousLabel = label - } - return (time: string) => { - return realtimeLabels.get(time) ?? formatHourMinute(time) - } - } - if (range.hours >= 168) { - return (time: string) => { - const d = new Date(time) - const mm = String(d.getMonth() + 1).padStart(2, '0') - const dd = String(d.getDate()).padStart(2, '0') - return `${mm}-${dd}` - } - } - return undefined - }, [isRealtime, chartData, range]) + return useMemo(() => makeTickFormatter(isRealtime, range.hours, chartData), [isRealtime, chartData, range]) } -/** Recharts X-axis `interval` (stride = N+1 ticks). Returns a numeric stride for - * long ranges so the auto-generated tick labels do not overlap on narrow viewports. - * 7d/30d hourly buckets produce 168 / 720 samples — labelling every one collapses - * into illegible overlap, so we target roughly 8 evenly spaced labels. */ function useXAxisInterval(isRealtime: boolean, range: TimeRange, dataLength: number) { - return useMemo(() => { - if (isRealtime) { - return 0 - } - if (range.hours >= 168 && dataLength > 0) { - const targetLabels = 8 - return Math.max(0, Math.floor(dataLength / targetLabels) - 1) - } - return undefined - }, [isRealtime, range, dataLength]) + return useMemo(() => xAxisStride(isRealtime, range.hours, dataLength), [isRealtime, range, dataLength]) } function useTooltipFormatter(isRealtime: boolean, range: TimeRange) { - return useMemo<((time: string) => string) | undefined>(() => { - if (isRealtime) { - return (time: string) => { - const d = new Date(time) - return `${String(d.getHours()).padStart(2, '0')}:${String(d.getMinutes()).padStart(2, '0')}:${String(d.getSeconds()).padStart(2, '0')}` - } - } - if (range.hours >= 168) { - return (time: string) => { - const d = new Date(time) - return `${String(d.getMonth() + 1).padStart(2, '0')}-${String(d.getDate()).padStart(2, '0')} ${String(d.getHours()).padStart(2, '0')}:${String(d.getMinutes()).padStart(2, '0')}` - } - } - return undefined - }, [isRealtime, range]) + return useMemo(() => makeTooltipFormatter(isRealtime, range.hours), [isRealtime, range]) } function useGpuChartData( @@ -491,23 +362,10 @@ function useGpuChartData( gpuRecords: GpuRecordAggregated[] | undefined, publicMetrics: PublicMetricsPoint[] | undefined ) { - return useMemo[]>(() => { - if (isAdminVariant) { - if (!gpuRecords || gpuRecords.length === 0) { - return [] - } - return gpuRecords.map((r) => ({ - timestamp: r.time, - gpu_usage: r.gpu_usage_avg, - gpu_temp: r.temperature_avg, - gpu_mem_pct: r.mem_total_avg > 0 ? (r.mem_used_avg / r.mem_total_avg) * 100 : 0 - })) - } - if (!publicMetrics) { - return [] - } - return publicMetrics.filter((p) => p.gpu_usage != null).map((p) => ({ timestamp: p.time, gpu_usage: p.gpu_usage })) - }, [isAdminVariant, gpuRecords, publicMetrics]) + return useMemo( + () => buildGpuChartRows(isAdminVariant, gpuRecords, publicMetrics), + [isAdminVariant, gpuRecords, publicMetrics] + ) } interface AvailableMetrics { diff --git a/apps/web/src/hooks/use-realtime-metrics.ts b/apps/web/src/hooks/use-realtime-metrics.ts index e0f50ad1f..ec1853e0e 100644 --- a/apps/web/src/hooks/use-realtime-metrics.ts +++ b/apps/web/src/hooks/use-realtime-metrics.ts @@ -1,6 +1,7 @@ import type { QueryClient } from '@tanstack/react-query' import { useQueryClient } from '@tanstack/react-query' import { useEffect, useRef, useState } from 'react' +import { pct } from '@/lib/metric-chart-model' import { readLiveServers, type ServerMetrics, subscribeLiveServers } from '@/lib/server-catalog' const MAX_BUFFER_SIZE = 200 @@ -47,11 +48,11 @@ export interface RealtimeDataPoint { export function toRealtimeDataPoint(metrics: ServerMetrics): RealtimeDataPoint { return { cpu: metrics.cpu, - disk_pct: metrics.disk_total > 0 ? (metrics.disk_used / metrics.disk_total) * 100 : 0, + disk_pct: pct(metrics.disk_used, metrics.disk_total), load1: metrics.load1, load5: metrics.load5, load15: metrics.load15, - memory_pct: metrics.mem_total > 0 ? (metrics.mem_used / metrics.mem_total) * 100 : 0, + memory_pct: pct(metrics.mem_used, metrics.mem_total), net_in_speed: metrics.net_in_speed, net_in_transfer: metrics.net_in_transfer, net_out_speed: metrics.net_out_speed, diff --git a/apps/web/src/lib/metric-chart-model.test.ts b/apps/web/src/lib/metric-chart-model.test.ts new file mode 100644 index 000000000..339ff2db5 --- /dev/null +++ b/apps/web/src/lib/metric-chart-model.test.ts @@ -0,0 +1,209 @@ +import { describe, expect, it } from 'vitest' +import type { PublicMetricsPoint } from '@/lib/api-schema' +import { + buildGpuChartRows, + buildRealtimeTickLabels, + deriveNetworkLabels, + type MetricSeriesPoint, + makeTickFormatter, + makeTooltipFormatter, + pct, + toMetricChartRow, + xAxisStride +} from './metric-chart-model' + +function makePoint(overrides: Partial = {}): MetricSeriesPoint { + return { + cpu: 42, + disk_used: 250, + load1: 1, + load5: 0.8, + load15: 0.6, + mem_used: 512, + net_in_speed: 1000, + net_in_transfer: 10_000, + net_out_speed: 500, + net_out_transfer: 5000, + temperature: 55, + time: '2026-07-11T10:00:00Z', + ...overrides + } +} + +function makePublicPoint(overrides: Partial = {}): PublicMetricsPoint { + return { + ...makePoint(), + gpu_usage: null, + process_count: 100, + tcp_conn: 10, + temperature: 55, + udp_conn: 2, + ...overrides + } +} + +describe('pct', () => { + it('computes a percentage against the total', () => { + expect(pct(512, 1024)).toBe(50) + }) + + it('guards zero and negative totals to 0 instead of NaN/Infinity', () => { + expect(pct(512, 0)).toBe(0) + expect(pct(512, -1)).toBe(0) + expect(pct(0, 0)).toBe(0) + }) +}) + +describe('toMetricChartRow', () => { + it('maps a series point into a chart row with derived percentages', () => { + const row = toMetricChartRow(makePoint(), 1024, 1000) + expect(row.timestamp).toBe('2026-07-11T10:00:00Z') + expect(row.cpu).toBe(42) + expect(row.memory_pct).toBe(50) + expect(row.disk_pct).toBe(25) + expect(row.net_in_speed).toBe(1000) + expect(row.load15).toBe(0.6) + expect(row.temperature).toBe(55) + }) + + it('renders 0% when the server row has no totals yet', () => { + const row = toMetricChartRow(makePoint(), 0, 0) + expect(row.memory_pct).toBe(0) + expect(row.disk_pct).toBe(0) + }) +}) + +describe('buildRealtimeTickLabels', () => { + it('labels only the first point of each minute', () => { + const labels = buildRealtimeTickLabels([ + { timestamp: '2026-07-11T10:00:05Z' }, + { timestamp: '2026-07-11T10:00:35Z' }, + { timestamp: '2026-07-11T10:01:05Z' } + ]) + expect(labels.get('2026-07-11T10:00:05Z')).toMatch(/^\d{2}:\d{2}$/) + expect(labels.get('2026-07-11T10:00:35Z')).toBe('') + expect(labels.get('2026-07-11T10:01:05Z')).toMatch(/^\d{2}:\d{2}$/) + }) + + it('skips rows without a string timestamp', () => { + const labels = buildRealtimeTickLabels([{ timestamp: 42 }, {}]) + expect(labels.size).toBe(0) + }) +}) + +describe('makeTickFormatter', () => { + it('falls back to HH:MM for unknown realtime timestamps', () => { + const format = makeTickFormatter(true, 0, []) + expect(format?.('2026-07-11T10:00:00Z')).toMatch(/^\d{2}:\d{2}$/) + }) + + it('uses MM-DD labels for ranges of 7 days and longer', () => { + const format = makeTickFormatter(false, 168, []) + expect(format?.('2026-07-11T10:00:00Z')).toMatch(/^\d{2}-\d{2}$/) + }) + + it('leaves short historical ranges to the chart default', () => { + expect(makeTickFormatter(false, 24, [])).toBeUndefined() + }) +}) + +describe('makeTooltipFormatter', () => { + it('shows seconds in realtime mode', () => { + const format = makeTooltipFormatter(true, 0) + expect(format?.('2026-07-11T10:00:07Z')).toMatch(/^\d{2}:\d{2}:\d{2}$/) + }) + + it('shows date and time for long ranges', () => { + const format = makeTooltipFormatter(false, 720) + expect(format?.('2026-07-11T10:00:00Z')).toMatch(/^\d{2}-\d{2} \d{2}:\d{2}$/) + }) + + it('leaves short historical ranges to the chart default', () => { + expect(makeTooltipFormatter(false, 24)).toBeUndefined() + }) +}) + +describe('xAxisStride', () => { + it('labels every realtime point', () => { + expect(xAxisStride(true, 0, 500)).toBe(0) + }) + + it('targets roughly 8 labels on long ranges', () => { + expect(xAxisStride(false, 168, 168)).toBe(20) + expect(xAxisStride(false, 720, 720)).toBe(89) + }) + + it('never returns a negative stride for tiny datasets', () => { + expect(xAxisStride(false, 168, 3)).toBe(0) + }) + + it('leaves short ranges and empty data to the chart default', () => { + expect(xAxisStride(false, 24, 100)).toBeUndefined() + expect(xAxisStride(false, 168, 0)).toBeUndefined() + }) +}) + +describe('buildGpuChartRows', () => { + it('derives gpu memory percentage with the zero-total guard', () => { + const rows = buildGpuChartRows( + true, + [ + { gpu_usage_avg: 80, mem_total_avg: 16, mem_used_avg: 8, temperature_avg: 70, time: 't1' }, + { gpu_usage_avg: 10, mem_total_avg: 0, mem_used_avg: 0, temperature_avg: 40, time: 't2' } + ], + undefined + ) + expect(rows).toEqual([ + { gpu_mem_pct: 50, gpu_temp: 70, gpu_usage: 80, timestamp: 't1' }, + { gpu_mem_pct: 0, gpu_temp: 40, gpu_usage: 10, timestamp: 't2' } + ]) + }) + + it('filters public points without gpu data', () => { + const rows = buildGpuChartRows(false, undefined, [ + makePublicPoint({ gpu_usage: 33, time: 'with-gpu' }), + makePublicPoint({ gpu_usage: null, time: 'without-gpu' }) + ]) + expect(rows).toEqual([{ gpu_usage: 33, timestamp: 'with-gpu' }]) + }) + + it('returns empty rows when no source data exists', () => { + expect(buildGpuChartRows(true, undefined, undefined)).toEqual([]) + expect(buildGpuChartRows(true, [], undefined)).toEqual([]) + expect(buildGpuChartRows(false, undefined, undefined)).toEqual([]) + }) +}) + +describe('deriveNetworkLabels', () => { + const snapshot = { + net_in_transfer: 1024, + net_out_transfer: 2048 + } as Parameters[2] + + it('renders placeholders while admin live data is missing', () => { + expect(deriveNetworkLabels(true, undefined, null)).toEqual({ + netInLabel: '—', + netOutLabel: '—', + netTotalLabel: '—' + }) + }) + + it('sums cumulative transfer for the admin total', () => { + const live = { net_in_transfer: 1024, net_out_transfer: 2048 } as Parameters[1] + const labels = deriveNetworkLabels(true, live, null) + expect(labels.netInLabel).toBe('1.0 KB') + expect(labels.netOutLabel).toBe('2.0 KB') + expect(labels.netTotalLabel).toBe('3.0 KB') + }) + + it('uses the public snapshot transfer totals, omitting the total when absent', () => { + expect(deriveNetworkLabels(false, undefined, null)).toEqual({ + netInLabel: '—', + netOutLabel: '—', + netTotalLabel: null + }) + const labels = deriveNetworkLabels(false, undefined, snapshot) + expect(labels.netInLabel).toBe('1.0 KB') + expect(labels.netTotalLabel).toBe('3.0 KB') + }) +}) diff --git a/apps/web/src/lib/metric-chart-model.ts b/apps/web/src/lib/metric-chart-model.ts new file mode 100644 index 000000000..ba8c46f6e --- /dev/null +++ b/apps/web/src/lib/metric-chart-model.ts @@ -0,0 +1,208 @@ +import type { PublicMetricsPoint, PublicServerDetail } from '@/lib/api-schema' +import type { ServerMetrics } from '@/lib/server-catalog' +import { formatBytes } from '@/lib/utils' + +/** + * Chart model for the server-detail Metrics tab: every transform between + * API/WS payloads and Recharts rows lives here, testable through this + * interface. Rendering (and the trivial admin/public branch glue) stays in + * `server-detail-content.tsx`. + */ + +/** Fields shared by the admin `ServerMetricRecord` and public `PublicMetricsPoint` series. */ +export interface MetricSeriesPoint { + cpu: number + disk_used: number + load1: number + load5: number + load15: number + mem_used: number + net_in_speed: number + net_in_transfer: number + net_out_speed: number + net_out_transfer: number + temperature?: number | null + time: string +} + +export type MetricChartRow = { + cpu: number + disk_pct: number + load1: number + load5: number + load15: number + memory_pct: number + net_in_speed: number + net_in_transfer: number + net_out_speed: number + net_out_transfer: number + temperature?: number | null + timestamp: string +} + +/** The one percentage guard: a missing/zero total renders as 0%, never NaN/Infinity. */ +export function pct(used: number, total: number): number { + return total > 0 ? (used / total) * 100 : 0 +} + +export function toMetricChartRow(p: MetricSeriesPoint, memTotal: number, diskTotal: number): MetricChartRow { + return { + cpu: p.cpu, + disk_pct: pct(p.disk_used, diskTotal), + load1: p.load1, + load5: p.load5, + load15: p.load15, + memory_pct: pct(p.mem_used, memTotal), + net_in_speed: p.net_in_speed, + net_in_transfer: p.net_in_transfer, + net_out_speed: p.net_out_speed, + net_out_transfer: p.net_out_transfer, + temperature: p.temperature, + timestamp: p.time + } +} + +export interface GpuRecordAggregated { + gpu_usage_avg: number + mem_total_avg: number + mem_used_avg: number + temperature_avg: number + time: string +} + +export function buildGpuChartRows( + isAdminVariant: boolean, + gpuRecords: GpuRecordAggregated[] | undefined, + publicMetrics: PublicMetricsPoint[] | undefined +): Record[] { + if (isAdminVariant) { + if (!gpuRecords || gpuRecords.length === 0) { + return [] + } + return gpuRecords.map((r) => ({ + gpu_mem_pct: pct(r.mem_used_avg, r.mem_total_avg), + gpu_temp: r.temperature_avg, + gpu_usage: r.gpu_usage_avg, + timestamp: r.time + })) + } + if (!publicMetrics) { + return [] + } + return publicMetrics.filter((p) => p.gpu_usage != null).map((p) => ({ gpu_usage: p.gpu_usage, timestamp: p.time })) +} + +function formatHourMinute(time: string) { + const d = new Date(time) + return `${String(d.getHours()).padStart(2, '0')}:${String(d.getMinutes()).padStart(2, '0')}` +} + +function formatMonthDay(time: string) { + const d = new Date(time) + return `${String(d.getMonth() + 1).padStart(2, '0')}-${String(d.getDate()).padStart(2, '0')}` +} + +/** Range length (hours) at which tick labels switch from HH:MM to MM-DD. */ +const LONG_RANGE_HOURS = 168 + +/** + * Realtime points arrive seconds apart, so consecutive ticks often share an + * HH:MM label — repeats collapse to '' and only the first of each minute is + * labelled. + */ +export function buildRealtimeTickLabels(rows: { timestamp?: unknown }[]): Map { + const labels = new Map() + let previousLabel = '' + for (const point of rows) { + if (typeof point.timestamp !== 'string') { + continue + } + const label = formatHourMinute(point.timestamp) + labels.set(point.timestamp, label === previousLabel ? '' : label) + previousLabel = label + } + return labels +} + +export function makeTickFormatter( + isRealtime: boolean, + rangeHours: number, + rows: { timestamp?: unknown }[] +): ((time: string) => string) | undefined { + if (isRealtime) { + const realtimeLabels = buildRealtimeTickLabels(rows) + return (time: string) => realtimeLabels.get(time) ?? formatHourMinute(time) + } + if (rangeHours >= LONG_RANGE_HOURS) { + return formatMonthDay + } + return undefined +} + +export function makeTooltipFormatter(isRealtime: boolean, rangeHours: number): ((time: string) => string) | undefined { + if (isRealtime) { + return (time: string) => { + const d = new Date(time) + return `${String(d.getHours()).padStart(2, '0')}:${String(d.getMinutes()).padStart(2, '0')}:${String(d.getSeconds()).padStart(2, '0')}` + } + } + if (rangeHours >= LONG_RANGE_HOURS) { + return (time: string) => `${formatMonthDay(time)} ${formatHourMinute(time)}` + } + return undefined +} + +/** + * Recharts X-axis `interval` (stride = N+1 ticks). Long ranges produce 168/720 + * hourly samples — labelling every one collapses into illegible overlap, so we + * target roughly 8 evenly spaced labels. + */ +export function xAxisStride(isRealtime: boolean, rangeHours: number, dataLength: number): number | undefined { + if (isRealtime) { + return 0 + } + if (rangeHours >= LONG_RANGE_HOURS && dataLength > 0) { + const targetLabels = 8 + return Math.max(0, Math.floor(dataLength / targetLabels) - 1) + } + return undefined +} + +export interface NetworkLabels { + netInLabel: string + netOutLabel: string + netTotalLabel: string | null +} + +export function deriveNetworkLabels( + isAdminVariant: boolean, + liveData: ServerMetrics | undefined, + publicMetricsSnapshot: PublicServerDetail['metrics'] | null +): NetworkLabels { + if (isAdminVariant) { + if (!liveData) { + return { netInLabel: '—', netOutLabel: '—', netTotalLabel: '—' } + } + const inBytes = liveData.net_in_transfer ?? 0 + const outBytes = liveData.net_out_transfer ?? 0 + return { + netInLabel: formatBytes(inBytes), + netOutLabel: formatBytes(outBytes), + netTotalLabel: formatBytes(inBytes + outBytes) + } + } + if (!publicMetricsSnapshot) { + return { netInLabel: '—', netOutLabel: '—', netTotalLabel: null } + } + // Use cumulative transfer (not the instantaneous *_speed rate) so the public + // bar matches the admin bar: the `detail_network_in/out/total` labels describe a + // total amount transferred, and formatBytes renders bytes — feeding a rate here + // mislabelled "1.2 MB/s" as a cumulative "1.2 MB". + const inBytes = publicMetricsSnapshot.net_in_transfer + const outBytes = publicMetricsSnapshot.net_out_transfer + return { + netInLabel: formatBytes(inBytes), + netOutLabel: formatBytes(outBytes), + netTotalLabel: formatBytes(inBytes + outBytes) + } +} From bc66a5128e9b90fcbabc0c94675cc99602be24a3 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 22:45:03 +0800 Subject: [PATCH 25/44] refactor(web): share one isoWindow helper across range queries Four hooks hand-rolled the same ISO {from, to} window construction (server records, network records, network anomalies, GPU records / public metrics). Extract isoWindow(hours) into lib/utils so the query window shape has one owner. --- .../components/status/server-detail-content.tsx | 14 +++----------- apps/web/src/hooks/use-api.ts | 5 ++--- apps/web/src/hooks/use-network-api.ts | 9 +++------ apps/web/src/lib/utils.ts | 10 ++++++++++ 4 files changed, 18 insertions(+), 20 deletions(-) diff --git a/apps/web/src/components/status/server-detail-content.tsx b/apps/web/src/components/status/server-detail-content.tsx index 6e0273775..8f9de6456 100644 --- a/apps/web/src/components/status/server-detail-content.tsx +++ b/apps/web/src/components/status/server-detail-content.tsx @@ -36,7 +36,7 @@ import { } from '@/lib/metric-chart-model' import { useLiveServers } from '@/lib/server-catalog' import { type RangeKey, rangesForVariant, resolveRange, type TimeRange } from '@/lib/server-detail-nav' -import { cn, formatBytes } from '@/lib/utils' +import { cn, formatBytes, isoWindow } from '@/lib/utils' import { computeAggregateUptime } from '@/lib/widget-helpers' export interface ServerDetailContentProps { @@ -66,14 +66,6 @@ function isAdminServer(server: ServerResponse | PublicServerDetail): server is S return 'ipv4' in server } -function buildIsoWindow(hours: number) { - const now = new Date() - return { - from: new Date(now.getTime() - hours * 3600 * 1000).toISOString(), - to: now.toISOString() - } -} - // Fetches the historical metric series, branching on variant. Admin uses the // auth'd `useServerRecords` (includes disk-io + temperature blobs); public // hits `/api/status/servers/{id}/metrics` which returns the normalised @@ -85,7 +77,7 @@ function useMetricSeries(serverId: string, range: TimeRange, isAdminVariant: boo const { data: publicMetrics } = useQuery({ queryKey: ['public-status', 'server', serverId, 'metrics', range.hours, range.interval], queryFn: () => { - const { from, to } = buildIsoWindow(range.hours) + const { from, to } = isoWindow(range.hours) return api.get( `/api/status/servers/${serverId}/metrics?from=${encodeURIComponent(from)}&to=${encodeURIComponent(to)}&interval=${encodeURIComponent(range.interval)}` ) @@ -100,7 +92,7 @@ function useAdminGpuRecords(serverId: string, range: TimeRange, isAdminVariant: return useQuery({ queryKey: ['servers', serverId, 'gpu-records', range.hours], queryFn: () => { - const { from, to } = buildIsoWindow(range.hours) + const { from, to } = isoWindow(range.hours) return api.get( `/api/servers/${serverId}/gpu-records?from=${encodeURIComponent(from)}&to=${encodeURIComponent(to)}` ) diff --git a/apps/web/src/hooks/use-api.ts b/apps/web/src/hooks/use-api.ts index 6433397a8..37c0146d8 100644 --- a/apps/web/src/hooks/use-api.ts +++ b/apps/web/src/hooks/use-api.ts @@ -2,6 +2,7 @@ import { useQuery } from '@tanstack/react-query' import { api } from '@/lib/api-client' import type { UptimeDailyEntry } from '@/lib/api-schema' import { useServerDetail } from '@/lib/server-catalog' +import { isoWindow } from '@/lib/utils' type ServerRecord = import('@/lib/api-schema').ServerMetricRecord @@ -13,9 +14,7 @@ export function useServerRecords(id: string, hours: number, interval: string, op return useQuery({ queryKey: ['servers', id, 'records', hours, interval], queryFn: () => { - const now = new Date() - const from = new Date(now.getTime() - hours * 3600 * 1000).toISOString() - const to = now.toISOString() + const { from, to } = isoWindow(hours) return api.get( `/api/servers/${id}/records?from=${encodeURIComponent(from)}&to=${encodeURIComponent(to)}&interval=${encodeURIComponent(interval)}` ) diff --git a/apps/web/src/hooks/use-network-api.ts b/apps/web/src/hooks/use-network-api.ts index 56d77f79a..82ccdf917 100644 --- a/apps/web/src/hooks/use-network-api.ts +++ b/apps/web/src/hooks/use-network-api.ts @@ -11,6 +11,7 @@ import type { TracerouteResponse, TracerouteResult } from '@/lib/network-types' +import { isoWindow } from '@/lib/utils' export function useNetworkTargets() { return useQuery({ @@ -46,9 +47,7 @@ export function useNetworkRecords(serverId: string, hours: number, options?: { t return useQuery({ queryKey: ['servers', serverId, 'network-probes', 'records', hours, options?.targetId], queryFn: () => { - const now = new Date() - const from = new Date(now.getTime() - hours * 3600 * 1000).toISOString() - const to = now.toISOString() + const { from, to } = isoWindow(hours) let url = `/api/servers/${serverId}/network-probes/records?from=${encodeURIComponent(from)}&to=${encodeURIComponent(to)}` if (options?.targetId) { url += `&target_id=${encodeURIComponent(options.targetId)}` @@ -64,9 +63,7 @@ export function useNetworkAnomalies(serverId: string, hours: number) { return useQuery({ queryKey: ['servers', serverId, 'network-probes', 'anomalies', hours], queryFn: () => { - const now = new Date() - const from = new Date(now.getTime() - hours * 3600 * 1000).toISOString() - const to = now.toISOString() + const { from, to } = isoWindow(hours) return api.get( `/api/servers/${serverId}/network-probes/anomalies?from=${encodeURIComponent(from)}&to=${encodeURIComponent(to)}` ) diff --git a/apps/web/src/lib/utils.ts b/apps/web/src/lib/utils.ts index 181018398..1a4536a3d 100644 --- a/apps/web/src/lib/utils.ts +++ b/apps/web/src/lib/utils.ts @@ -47,6 +47,16 @@ export function countryCodeToName(code: string | null | undefined, locale = 'en' } } +/** ISO `{from, to}` window ending now and reaching `hours` back — the query + * window shape every records/anomalies endpoint expects. */ +export function isoWindow(hours: number): { from: string; to: string } { + const now = new Date() + return { + from: new Date(now.getTime() - hours * 3600 * 1000).toISOString(), + to: now.toISOString() + } +} + export function formatBytes(bytes: number): string { if (bytes <= 0 || !Number.isFinite(bytes)) { return '0 B' From 4ec3477297ab70b34b0651ea13d3b4e19fcf3837 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 22:45:03 +0800 Subject: [PATCH 26/44] style(web): settle lint on chart model and ws test files Hoist tick/tooltip format regexes to module scope, use T[] shorthand, and document why MetricChartRow stays a type alias (interfaces lack an implicit index signature for Record chart props). --- apps/web/src/hooks/use-servers-ws.test.ts | 4 ++-- apps/web/src/lib/metric-chart-model.test.ts | 17 +++++++++++------ apps/web/src/lib/metric-chart-model.ts | 1 + 3 files changed, 14 insertions(+), 8 deletions(-) diff --git a/apps/web/src/hooks/use-servers-ws.test.ts b/apps/web/src/hooks/use-servers-ws.test.ts index 5513e22ef..d7aa7e62b 100644 --- a/apps/web/src/hooks/use-servers-ws.test.ts +++ b/apps/web/src/hooks/use-servers-ws.test.ts @@ -96,7 +96,7 @@ describe('handleWsMessage upgrade messages', () => { describe('handleWsMessage network probe updates', () => { it('dispatches validated network_probe_update frames to subscribers', () => { const queryClient = new QueryClient() - const received: Array> = [] + const received: Record[] = [] const unsubscribe = subscribeBrowserMessage('network_probe_update', (msg) => received.push(msg)) handleWsMessage( @@ -115,7 +115,7 @@ describe('handleWsMessage network probe updates', () => { it('drops malformed network_probe_update frames before dispatch', () => { const queryClient = new QueryClient() - const received: Array> = [] + const received: Record[] = [] const unsubscribe = subscribeBrowserMessage('network_probe_update', (msg) => received.push(msg)) handleWsMessage({ results: [null], server_id: 's1', type: 'network_probe_update' }, queryClient) diff --git a/apps/web/src/lib/metric-chart-model.test.ts b/apps/web/src/lib/metric-chart-model.test.ts index 339ff2db5..ae62527b2 100644 --- a/apps/web/src/lib/metric-chart-model.test.ts +++ b/apps/web/src/lib/metric-chart-model.test.ts @@ -12,6 +12,11 @@ import { xAxisStride } from './metric-chart-model' +const HH_MM = /^\d{2}:\d{2}$/ +const MM_DD = /^\d{2}-\d{2}$/ +const HH_MM_SS = /^\d{2}:\d{2}:\d{2}$/ +const MM_DD_HH_MM = /^\d{2}-\d{2} \d{2}:\d{2}$/ + function makePoint(overrides: Partial = {}): MetricSeriesPoint { return { cpu: 42, @@ -80,9 +85,9 @@ describe('buildRealtimeTickLabels', () => { { timestamp: '2026-07-11T10:00:35Z' }, { timestamp: '2026-07-11T10:01:05Z' } ]) - expect(labels.get('2026-07-11T10:00:05Z')).toMatch(/^\d{2}:\d{2}$/) + expect(labels.get('2026-07-11T10:00:05Z')).toMatch(HH_MM) expect(labels.get('2026-07-11T10:00:35Z')).toBe('') - expect(labels.get('2026-07-11T10:01:05Z')).toMatch(/^\d{2}:\d{2}$/) + expect(labels.get('2026-07-11T10:01:05Z')).toMatch(HH_MM) }) it('skips rows without a string timestamp', () => { @@ -94,12 +99,12 @@ describe('buildRealtimeTickLabels', () => { describe('makeTickFormatter', () => { it('falls back to HH:MM for unknown realtime timestamps', () => { const format = makeTickFormatter(true, 0, []) - expect(format?.('2026-07-11T10:00:00Z')).toMatch(/^\d{2}:\d{2}$/) + expect(format?.('2026-07-11T10:00:00Z')).toMatch(HH_MM) }) it('uses MM-DD labels for ranges of 7 days and longer', () => { const format = makeTickFormatter(false, 168, []) - expect(format?.('2026-07-11T10:00:00Z')).toMatch(/^\d{2}-\d{2}$/) + expect(format?.('2026-07-11T10:00:00Z')).toMatch(MM_DD) }) it('leaves short historical ranges to the chart default', () => { @@ -110,12 +115,12 @@ describe('makeTickFormatter', () => { describe('makeTooltipFormatter', () => { it('shows seconds in realtime mode', () => { const format = makeTooltipFormatter(true, 0) - expect(format?.('2026-07-11T10:00:07Z')).toMatch(/^\d{2}:\d{2}:\d{2}$/) + expect(format?.('2026-07-11T10:00:07Z')).toMatch(HH_MM_SS) }) it('shows date and time for long ranges', () => { const format = makeTooltipFormatter(false, 720) - expect(format?.('2026-07-11T10:00:00Z')).toMatch(/^\d{2}-\d{2} \d{2}:\d{2}$/) + expect(format?.('2026-07-11T10:00:00Z')).toMatch(MM_DD_HH_MM) }) it('leaves short historical ranges to the chart default', () => { diff --git a/apps/web/src/lib/metric-chart-model.ts b/apps/web/src/lib/metric-chart-model.ts index ba8c46f6e..16eb1fc97 100644 --- a/apps/web/src/lib/metric-chart-model.ts +++ b/apps/web/src/lib/metric-chart-model.ts @@ -25,6 +25,7 @@ export interface MetricSeriesPoint { time: string } +// biome-ignore lint/style/useConsistentTypeDefinitions: interfaces lack an implicit index signature, and chart rows must be assignable to the Record[] the chart components accept export type MetricChartRow = { cpu: number disk_pct: number From 98860a181269e4b2bdb1656d82de55127992f048 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 22:59:17 +0800 Subject: [PATCH 27/44] fix(web): use existing chart_gpu_usage key for the GPU usage chart title t('chart_gpu') has no entry in either locale, so the GPU usage chart rendered the raw key as its title. --- apps/web/src/components/status/server-detail-content.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/web/src/components/status/server-detail-content.tsx b/apps/web/src/components/status/server-detail-content.tsx index 8f9de6456..677336f7e 100644 --- a/apps/web/src/components/status/server-detail-content.tsx +++ b/apps/web/src/components/status/server-detail-content.tsx @@ -505,7 +505,7 @@ function MetricsTabContent({ domain={[0, 100]} formatTime={formatTime} formatTooltipLabel={formatTooltipLabel} - title={t('chart_gpu')} + title={t('chart_gpu_usage')} unit="%" xAxisInterval={xAxisInterval} /> From b799f1cea7187fd858a8c51d2bb5709e9006ded2 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 23:09:48 +0800 Subject: [PATCH 28/44] refactor(server): generate hourly rollup SQL from a metric column descriptor Each scalar column of records/records_hourly appeared three times in a hand-written SQL string (INSERT list, SELECT aggregate, ON CONFLICT update), and the AVG-vs-MAX aggregation decision lived only inside that string. Declare every column once in service::rollup::METRIC_COLUMNS with its RollupAgg and generate the upsert from the table. A golden test pins the generated statement byte-for-byte to the original SQL. --- crates/server/src/service/mod.rs | 1 + crates/server/src/service/record.rs | 49 +-------- crates/server/src/service/rollup.rs | 154 ++++++++++++++++++++++++++++ 3 files changed, 159 insertions(+), 45 deletions(-) create mode 100644 crates/server/src/service/rollup.rs diff --git a/crates/server/src/service/mod.rs b/crates/server/src/service/mod.rs index b24fef922..253151090 100644 --- a/crates/server/src/service/mod.rs +++ b/crates/server/src/service/mod.rs @@ -30,6 +30,7 @@ pub mod oauth; pub mod ping; pub mod public_status; pub mod record; +pub mod rollup; pub mod security; pub mod server; pub mod server_tag; diff --git a/crates/server/src/service/record.rs b/crates/server/src/service/record.rs index 9e0fe91e5..aecc9f45b 100644 --- a/crates/server/src/service/record.rs +++ b/crates/server/src/service/record.rs @@ -211,51 +211,10 @@ impl RecordService { let hour_start_str = hour_start.to_rfc3339_opts(SecondsFormat::AutoSi, false); let hour_end_str = hour_end.to_rfc3339_opts(SecondsFormat::AutoSi, false); - // SQL aggregation for numeric columns with upsert - let sql = "INSERT INTO records_hourly \ - (server_id, time, cpu, mem_used, swap_used, disk_used, \ - net_in_speed, net_out_speed, net_in_transfer, net_out_transfer, \ - load1, load5, load15, tcp_conn, udp_conn, process_count, \ - temperature, gpu_usage) \ - SELECT \ - server_id, \ - ?, \ - AVG(cpu), \ - CAST(AVG(mem_used) AS INTEGER), \ - CAST(AVG(swap_used) AS INTEGER), \ - CAST(AVG(disk_used) AS INTEGER), \ - CAST(AVG(net_in_speed) AS INTEGER), \ - CAST(AVG(net_out_speed) AS INTEGER), \ - CAST(MAX(net_in_transfer) AS INTEGER), \ - CAST(MAX(net_out_transfer) AS INTEGER), \ - AVG(load1), \ - AVG(load5), \ - AVG(load15), \ - CAST(AVG(tcp_conn) AS INTEGER), \ - CAST(AVG(udp_conn) AS INTEGER), \ - CAST(AVG(process_count) AS INTEGER), \ - AVG(temperature), \ - AVG(gpu_usage) \ - FROM records \ - WHERE time >= ? AND time < ? \ - GROUP BY server_id \ - ON CONFLICT(server_id, time) DO UPDATE SET \ - cpu = excluded.cpu, \ - mem_used = excluded.mem_used, \ - swap_used = excluded.swap_used, \ - disk_used = excluded.disk_used, \ - net_in_speed = excluded.net_in_speed, \ - net_out_speed = excluded.net_out_speed, \ - net_in_transfer = excluded.net_in_transfer, \ - net_out_transfer = excluded.net_out_transfer, \ - load1 = excluded.load1, \ - load5 = excluded.load5, \ - load15 = excluded.load15, \ - tcp_conn = excluded.tcp_conn, \ - udp_conn = excluded.udp_conn, \ - process_count = excluded.process_count, \ - temperature = excluded.temperature, \ - gpu_usage = excluded.gpu_usage"; + // Scalar columns roll up in SQL; the statement is generated from the + // rollup-policy descriptor so column list and aggregation choices have + // a single owner (see service::rollup). + let sql = super::rollup::aggregate_hourly_sql(); let result = db .execute(Statement::from_sql_and_values( diff --git a/crates/server/src/service/rollup.rs b/crates/server/src/service/rollup.rs new file mode 100644 index 000000000..5625e738c --- /dev/null +++ b/crates/server/src/service/rollup.rs @@ -0,0 +1,154 @@ +//! Rollup policy: the single owner of how raw metric records compress over +//! time. Every scalar column of `records` / `records_hourly` is declared once +//! in [`METRIC_COLUMNS`] together with its aggregation function; the hourly +//! rollup SQL is generated from that table, so adding a metric column means +//! adding one descriptor row instead of hand-editing column lists in three +//! places of an SQL string. +//! +//! `disk_io_json` is the documented exception: SQLite cannot aggregate the +//! per-device JSON blob, so `RecordService::aggregate_hourly` folds it in Rust +//! after the SQL pass. + +/// How a raw metric column folds into its hourly bucket. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RollupAgg { + /// `AVG(col)` — REAL gauges keep fractional precision. + Avg, + /// `CAST(AVG(col) AS INTEGER)` — integer gauges. + AvgInt, + /// `CAST(MAX(col) AS INTEGER)` — cumulative counters keep the window-end + /// value instead of a meaningless average. + MaxInt, +} + +impl RollupAgg { + fn select_expr(self, column: &str) -> String { + match self { + RollupAgg::Avg => format!("AVG({column})"), + RollupAgg::AvgInt => format!("CAST(AVG({column}) AS INTEGER)"), + RollupAgg::MaxInt => format!("CAST(MAX({column}) AS INTEGER)"), + } + } +} + +/// One rollup-managed metric column, shared by `records` and `records_hourly` +/// (the two tables have identical column sets). +pub struct MetricColumn { + /// SQL column name in both tables. + pub name: &'static str, + /// How the raw samples fold into the hourly bucket. + pub agg: RollupAgg, +} + +/// The scalar metric columns, in physical column order. Order matters: the +/// generated INSERT lists columns in this order. +pub const METRIC_COLUMNS: &[MetricColumn] = &[ + MetricColumn { name: "cpu", agg: RollupAgg::Avg }, + MetricColumn { name: "mem_used", agg: RollupAgg::AvgInt }, + MetricColumn { name: "swap_used", agg: RollupAgg::AvgInt }, + MetricColumn { name: "disk_used", agg: RollupAgg::AvgInt }, + MetricColumn { name: "net_in_speed", agg: RollupAgg::AvgInt }, + MetricColumn { name: "net_out_speed", agg: RollupAgg::AvgInt }, + MetricColumn { name: "net_in_transfer", agg: RollupAgg::MaxInt }, + MetricColumn { name: "net_out_transfer", agg: RollupAgg::MaxInt }, + MetricColumn { name: "load1", agg: RollupAgg::Avg }, + MetricColumn { name: "load5", agg: RollupAgg::Avg }, + MetricColumn { name: "load15", agg: RollupAgg::Avg }, + MetricColumn { name: "tcp_conn", agg: RollupAgg::AvgInt }, + MetricColumn { name: "udp_conn", agg: RollupAgg::AvgInt }, + MetricColumn { name: "process_count", agg: RollupAgg::AvgInt }, + MetricColumn { name: "temperature", agg: RollupAgg::Avg }, + MetricColumn { name: "gpu_usage", agg: RollupAgg::Avg }, +]; + +/// Build the hourly rollup upsert. Placeholders: bucket time, window start, +/// window end (in that order). +pub fn aggregate_hourly_sql() -> String { + let insert_columns = METRIC_COLUMNS + .iter() + .map(|c| c.name) + .collect::>() + .join(", "); + let select_exprs = METRIC_COLUMNS + .iter() + .map(|c| c.agg.select_expr(c.name)) + .collect::>() + .join(", "); + let conflict_sets = METRIC_COLUMNS + .iter() + .map(|c| format!("{0} = excluded.{0}", c.name)) + .collect::>() + .join(", "); + format!( + "INSERT INTO records_hourly (server_id, time, {insert_columns}) \ + SELECT server_id, ?, {select_exprs} \ + FROM records WHERE time >= ? AND time < ? \ + GROUP BY server_id \ + ON CONFLICT(server_id, time) DO UPDATE SET {conflict_sets}" + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Golden test: the generated statement must stay byte-for-byte identical + /// to the hand-written SQL it replaced, so the descriptor migration cannot + /// change rollup semantics. + #[test] + fn aggregate_hourly_sql_matches_the_original_hand_written_statement() { + let expected = "INSERT INTO records_hourly \ + (server_id, time, cpu, mem_used, swap_used, disk_used, \ + net_in_speed, net_out_speed, net_in_transfer, net_out_transfer, \ + load1, load5, load15, tcp_conn, udp_conn, process_count, \ + temperature, gpu_usage) \ + SELECT \ + server_id, \ + ?, \ + AVG(cpu), \ + CAST(AVG(mem_used) AS INTEGER), \ + CAST(AVG(swap_used) AS INTEGER), \ + CAST(AVG(disk_used) AS INTEGER), \ + CAST(AVG(net_in_speed) AS INTEGER), \ + CAST(AVG(net_out_speed) AS INTEGER), \ + CAST(MAX(net_in_transfer) AS INTEGER), \ + CAST(MAX(net_out_transfer) AS INTEGER), \ + AVG(load1), \ + AVG(load5), \ + AVG(load15), \ + CAST(AVG(tcp_conn) AS INTEGER), \ + CAST(AVG(udp_conn) AS INTEGER), \ + CAST(AVG(process_count) AS INTEGER), \ + AVG(temperature), \ + AVG(gpu_usage) \ + FROM records \ + WHERE time >= ? AND time < ? \ + GROUP BY server_id \ + ON CONFLICT(server_id, time) DO UPDATE SET \ + cpu = excluded.cpu, \ + mem_used = excluded.mem_used, \ + swap_used = excluded.swap_used, \ + disk_used = excluded.disk_used, \ + net_in_speed = excluded.net_in_speed, \ + net_out_speed = excluded.net_out_speed, \ + net_in_transfer = excluded.net_in_transfer, \ + net_out_transfer = excluded.net_out_transfer, \ + load1 = excluded.load1, \ + load5 = excluded.load5, \ + load15 = excluded.load15, \ + tcp_conn = excluded.tcp_conn, \ + udp_conn = excluded.udp_conn, \ + process_count = excluded.process_count, \ + temperature = excluded.temperature, \ + gpu_usage = excluded.gpu_usage"; + assert_eq!(aggregate_hourly_sql(), expected); + } + + #[test] + fn metric_columns_have_unique_names() { + let mut names: Vec<_> = METRIC_COLUMNS.iter().map(|c| c.name).collect(); + names.sort_unstable(); + names.dedup(); + assert_eq!(names.len(), METRIC_COLUMNS.len()); + } +} From 25a845222bbc719e5dd7271069288b7b83cbccb6 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 23:14:00 +0800 Subject: [PATCH 29/44] refactor(server): move raw/hourly table selection into the rollup policy query_history hardcoded the 24h auto-switch inline, and the unit test covering it re-implemented the branch instead of calling it. Own the switch point in service::rollup (select_history_table + the RAW_WINDOW_MAX_HOURS constant) so bucket choice and aggregation rules live in the same module, and test the real function. --- crates/server/src/service/record.rs | 55 ++--------------------------- crates/server/src/service/rollup.rs | 53 +++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 53 deletions(-) diff --git a/crates/server/src/service/record.rs b/crates/server/src/service/record.rs index aecc9f45b..2c26c0e12 100644 --- a/crates/server/src/service/record.rs +++ b/crates/server/src/service/record.rs @@ -147,17 +147,9 @@ impl RecordService { to: DateTime, interval: &str, ) -> Result { - let use_hourly = match interval { - "raw" => false, - "hourly" => true, - _ => { - // "auto" mode - let duration = to - from; - duration > Duration::hours(24) - } - }; + let table = super::rollup::select_history_table(interval, from, to); - if use_hourly { + if table == super::rollup::HistoryTable::Hourly { let records = record_hourly::Entity::find() .filter(record_hourly::Column::ServerId.eq(server_id)) .filter(record_hourly::Column::Time.gte(from)) @@ -1246,49 +1238,6 @@ mod tests { assert!(matches!(hourly, QueryHistoryResult::Hourly(v) if v.is_empty())); } - /// The interval auto-selection logic: <=24h => raw, >24h => hourly. - /// Extracted from `query_history` so we can verify it without DB. - #[test] - fn test_interval_selection_logic() { - let now = Utc::now(); - - // Within 24 hours => should use raw - let from_recent = now - Duration::hours(12); - let duration_recent = now - from_recent; - let use_hourly_recent = duration_recent > Duration::hours(24); - assert!(!use_hourly_recent, "12h range should select raw"); - - // Exactly 24 hours => should use raw (not >24h) - let from_exact = now - Duration::hours(24); - let duration_exact = now - from_exact; - let use_hourly_exact = duration_exact > Duration::hours(24); - assert!( - !use_hourly_exact, - "24h range should select raw (not strictly greater)" - ); - - // More than 24 hours => should use hourly - let from_old = now - Duration::hours(48); - let duration_old = now - from_old; - let use_hourly_old = duration_old > Duration::hours(24); - assert!(use_hourly_old, "48h range should select hourly"); - - // Explicit interval overrides - let explicit_raw = match "raw" { - "raw" => false, - "hourly" => true, - _ => unreachable!(), - }; - assert!(!explicit_raw, "explicit 'raw' should use raw"); - - let explicit_hourly = match "hourly" { - "raw" => false, - "hourly" => true, - _ => unreachable!(), - }; - assert!(explicit_hourly, "explicit 'hourly' should use hourly"); - } - /// Verify the retention cutoff calculation used in cleanup_expired. #[test] fn test_retention_cutoff_calculation() { diff --git a/crates/server/src/service/rollup.rs b/crates/server/src/service/rollup.rs index 5625e738c..7d2d5dedf 100644 --- a/crates/server/src/service/rollup.rs +++ b/crates/server/src/service/rollup.rs @@ -9,6 +9,36 @@ //! per-device JSON blob, so `RecordService::aggregate_hourly` folds it in Rust //! after the SQL pass. +use chrono::{DateTime, Duration, Utc}; + +/// Longest range the raw table serves under `interval = "auto"`; longer +/// ranges read the hourly rollup. +pub const RAW_WINDOW_MAX_HOURS: i64 = 24; + +/// Which table serves a history query. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum HistoryTable { + Raw, + Hourly, +} + +/// Resolve the table for a history query: explicit `"raw"` / `"hourly"` win, +/// anything else (the `"auto"` default) switches to the hourly rollup once +/// the range exceeds [`RAW_WINDOW_MAX_HOURS`]. +pub fn select_history_table(interval: &str, from: DateTime, to: DateTime) -> HistoryTable { + match interval { + "raw" => HistoryTable::Raw, + "hourly" => HistoryTable::Hourly, + _ => { + if to - from > Duration::hours(RAW_WINDOW_MAX_HOURS) { + HistoryTable::Hourly + } else { + HistoryTable::Raw + } + } + } +} + /// How a raw metric column folds into its hourly bucket. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum RollupAgg { @@ -144,6 +174,29 @@ mod tests { assert_eq!(aggregate_hourly_sql(), expected); } + /// Table selection: explicit intervals win; "auto" switches to hourly + /// strictly beyond 24h (a range of exactly 24h still reads raw). + #[test] + fn select_history_table_covers_explicit_and_auto_intervals() { + let now = Utc::now(); + let hours = |h: i64| now - Duration::hours(h); + + assert_eq!(select_history_table("raw", hours(48), now), HistoryTable::Raw); + assert_eq!( + select_history_table("hourly", hours(1), now), + HistoryTable::Hourly + ); + assert_eq!(select_history_table("auto", hours(12), now), HistoryTable::Raw); + assert_eq!( + select_history_table("auto", hours(RAW_WINDOW_MAX_HOURS), now), + HistoryTable::Raw + ); + assert_eq!( + select_history_table("auto", hours(48), now), + HistoryTable::Hourly + ); + } + #[test] fn metric_columns_have_unique_names() { let mut names: Vec<_> = METRIC_COLUMNS.iter().map(|c| c.name).collect(); From 972e473a6153e4d4d192385ae521d03556b02145 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 23:25:23 +0800 Subject: [PATCH 30/44] refactor(server): route alert metric reads through the rollup descriptor Alert evaluation kept a second private read path: check_threshold and get_last_record_time queried the records table directly, and extract_metric re-encoded the rule_type-to-column mapping as a match that nothing cross-checked against storage. The descriptor now carries alert_rule_type plus a typed accessor per column (rollup::alert_metric), and alert reads go through RecordService::query_recent / latest_record_time, so a storage change cannot silently detach alerting from what the recorder writes. Transfer counters stay deliberately non-alertable (alert_rule_type: None), matching the old fallback. --- crates/server/src/service/alert.rs | 180 +++--------------------- crates/server/src/service/record.rs | 31 ++++ crates/server/src/service/rollup.rs | 210 +++++++++++++++++++++++++--- 3 files changed, 246 insertions(+), 175 deletions(-) diff --git a/crates/server/src/service/alert.rs b/crates/server/src/service/alert.rs index 158fc57f8..2dafc4cbf 100644 --- a/crates/server/src/service/alert.rs +++ b/crates/server/src/service/alert.rs @@ -5,11 +5,13 @@ use sea_orm::*; use serde::{Deserialize, Serialize}; use uuid::Uuid; -use crate::entity::{alert_rule, alert_state, network_probe_record, record, server}; +use crate::entity::{alert_rule, alert_state, network_probe_record, server}; use crate::error::AppError; use crate::service::agent_manager::AgentManager; use crate::service::maintenance::MaintenanceService; use crate::service::notification::{NotificationService, NotifyContext}; +use crate::service::record::RecordService; +use crate::service::rollup; /// Rule types that are event-driven (not evaluated on a polling interval). const EVENT_DRIVEN_RULE_TYPES: &[&str] = &[ @@ -761,7 +763,11 @@ impl AlertService { let duration = item.duration.unwrap_or(60) as u64; !agent_manager.is_online(server_id) && { // Check how long the server has been offline by looking at last record - match get_last_record_time(db, server_id).await { + match RecordService::latest_record_time(db, server_id) + .await + .ok() + .flatten() + { Some(last) => { let elapsed = (Utc::now() - last).num_seconds().max(0) as u64; elapsed >= duration @@ -1011,30 +1017,8 @@ async fn resolve_servers( } } -async fn get_last_record_time( - db: &DatabaseConnection, - server_id: &str, -) -> Option> { - record::Entity::find() - .filter(record::Column::ServerId.eq(server_id)) - .order_by_desc(record::Column::Time) - .one(db) - .await - .ok() - .flatten() - .map(|r| r.time) -} - async fn check_threshold(db: &DatabaseConnection, server_id: &str, item: &AlertRuleItem) -> bool { - let ten_min_ago = Utc::now() - Duration::minutes(10); - - let records = match record::Entity::find() - .filter(record::Column::ServerId.eq(server_id)) - .filter(record::Column::Time.gte(ten_min_ago)) - .order_by_desc(record::Column::Time) - .all(db) - .await - { + let records = match RecordService::query_recent(db, server_id, Duration::minutes(10)).await { Ok(r) => r, Err(_) => return false, }; @@ -1047,7 +1031,7 @@ async fn check_threshold(db: &DatabaseConnection, server_id: &str, item: &AlertR let total = records.len(); for rec in &records { - let value = extract_metric(rec, &item.rule_type); + let value = rollup::alert_metric(rec, &item.rule_type); let exceeds = match (item.min, item.max) { (Some(min), Some(max)) => value >= min && value <= max, (Some(min), None) => value >= min, @@ -1232,30 +1216,6 @@ async fn check_network_packet_loss( } } -fn extract_metric(rec: &record::Model, rule_type: &str) -> f64 { - match rule_type { - "cpu" => rec.cpu, - "memory" => { - // We need mem_total for percentage but we don't have it in the record. - // Return raw mem_used as bytes. The threshold should be set accordingly. - rec.mem_used as f64 - } - "swap" => rec.swap_used as f64, - "disk" => rec.disk_used as f64, - "load1" => rec.load1, - "load5" => rec.load5, - "load15" => rec.load15, - "tcp_conn" => rec.tcp_conn as f64, - "udp_conn" => rec.udp_conn as f64, - "process" => rec.process_count as f64, - "net_in_speed" => rec.net_in_speed as f64, - "net_out_speed" => rec.net_out_speed as f64, - "temperature" => rec.temperature.unwrap_or(0.0), - "gpu" => rec.gpu_usage.unwrap_or(0.0), - _ => 0.0, - } -} - /// Render a compact human summary of a rule's conditions for notification /// bodies, e.g. `cpu >= 90 for 60s` or `offline for 120s; load1 >= 4`. /// Returns an empty string when `rules_json` cannot be parsed, so callers can @@ -1308,6 +1268,7 @@ fn majority_exceeded(exceeded_count: usize, total: usize) -> bool { #[cfg(test)] mod tests { use super::*; + use crate::entity::record; use chrono::{TimeZone, Utc}; // ── describe_rule_items: notification condition summary ── @@ -1330,32 +1291,6 @@ mod tests { assert_eq!(describe_rule_items("not json"), ""); } - /// Build a minimal `record::Model` with the given field values. - fn make_record(cpu: f64, mem_used: i64, load1: f64) -> record::Model { - record::Model { - id: 1, - server_id: "srv-1".to_string(), - time: Utc::now(), - cpu, - mem_used, - swap_used: 0, - disk_used: 0, - net_in_speed: 0, - net_out_speed: 0, - net_in_transfer: 0, - net_out_transfer: 0, - load1, - load5: 0.0, - load15: 0.0, - tcp_conn: 100, - udp_conn: 50, - process_count: 200, - temperature: Some(55.0), - gpu_usage: Some(40.0), - disk_io_json: None, - } - } - // ── T2-1: threshold above ── #[test] @@ -1448,55 +1383,6 @@ mod tests { assert!(majority_exceeded(1, 1), "1/1 = 100% should meet threshold"); } - // ── T2-4: extract_metric ── - - #[test] - fn test_extract_metric_cpu() { - let rec = make_record(85.5, 4_000_000, 1.2); - assert!((extract_metric(&rec, "cpu") - 85.5).abs() < f64::EPSILON); - } - - #[test] - fn test_extract_metric_memory() { - let rec = make_record(50.0, 8_000_000, 0.0); - assert!((extract_metric(&rec, "memory") - 8_000_000.0).abs() < f64::EPSILON); - } - - #[test] - fn test_extract_metric_load() { - let rec = make_record(0.0, 0, std::f64::consts::PI); - assert!((extract_metric(&rec, "load1") - std::f64::consts::PI).abs() < f64::EPSILON); - } - - #[test] - fn test_extract_metric_temperature() { - let rec = make_record(0.0, 0, 0.0); - assert!((extract_metric(&rec, "temperature") - 55.0).abs() < f64::EPSILON); - } - - #[test] - fn test_extract_metric_gpu() { - let rec = make_record(0.0, 0, 0.0); - assert!((extract_metric(&rec, "gpu") - 40.0).abs() < f64::EPSILON); - } - - #[test] - fn test_extract_metric_connections() { - let rec = make_record(0.0, 0, 0.0); - assert!((extract_metric(&rec, "tcp_conn") - 100.0).abs() < f64::EPSILON); - assert!((extract_metric(&rec, "udp_conn") - 50.0).abs() < f64::EPSILON); - assert!((extract_metric(&rec, "process") - 200.0).abs() < f64::EPSILON); - } - - #[test] - fn test_extract_metric_unknown_type() { - let rec = make_record(99.0, 0, 0.0); - assert!( - (extract_metric(&rec, "nonexistent") - 0.0).abs() < f64::EPSILON, - "unknown metric type should return 0.0" - ); - } - // ── T2-5: AlertRuleItem serialization round-trip ── #[test] @@ -2123,36 +2009,6 @@ mod tests { assert!(p.exclude_cidrs.is_empty()); } - // ── extract_metric: remaining variants ── - - #[test] - fn test_extract_metric_remaining_variants() { - // Exercise every metric branch not covered by the existing tests. - let mut rec = make_record(0.0, 0, 0.0); - rec.swap_used = 1024; - rec.disk_used = 2048; - rec.load5 = 2.5; - rec.load15 = 3.5; - rec.net_in_speed = 111; - rec.net_out_speed = 222; - assert!((extract_metric(&rec, "swap") - 1024.0).abs() < f64::EPSILON); - assert!((extract_metric(&rec, "disk") - 2048.0).abs() < f64::EPSILON); - assert!((extract_metric(&rec, "load5") - 2.5).abs() < f64::EPSILON); - assert!((extract_metric(&rec, "load15") - 3.5).abs() < f64::EPSILON); - assert!((extract_metric(&rec, "net_in_speed") - 111.0).abs() < f64::EPSILON); - assert!((extract_metric(&rec, "net_out_speed") - 222.0).abs() < f64::EPSILON); - } - - #[test] - fn test_extract_metric_temperature_and_gpu_none() { - // None temperature/gpu must coerce to 0.0 via unwrap_or. - let mut rec = make_record(0.0, 0, 0.0); - rec.temperature = None; - rec.gpu_usage = None; - assert!((extract_metric(&rec, "temperature") - 0.0).abs() < f64::EPSILON); - assert!((extract_metric(&rec, "gpu") - 0.0).abs() < f64::EPSILON); - } - // ── AlertRuleAction serde round-trip ── #[test] @@ -2624,13 +2480,16 @@ mod tests { assert_eq!(all.len(), 2); } - // ── get_last_record_time ── + // ── latest_record_time (shared RecordService reader) ── #[tokio::test] - async fn test_get_last_record_time_picks_latest_and_none() { + async fn test_latest_record_time_picks_latest_and_none() { let (db, _tmp) = setup_test_db().await; // No records => None. - assert!(get_last_record_time(&db, "srv-1").await.is_none()); + assert!(RecordService::latest_record_time(&db, "srv-1") + .await + .expect("query should succeed") + .is_none()); let older = Utc.with_ymd_and_hms(2026, 1, 1, 0, 0, 0).unwrap(); let newer = Utc.with_ymd_and_hms(2026, 1, 1, 1, 0, 0).unwrap(); @@ -2638,7 +2497,10 @@ mod tests { insert_record(&db, "srv-1", newer, 20.0).await; // Returns the most recent record's time. - let last = get_last_record_time(&db, "srv-1").await.expect("some"); + let last = RecordService::latest_record_time(&db, "srv-1") + .await + .expect("query should succeed") + .expect("some"); assert_eq!(last, newer); } diff --git a/crates/server/src/service/record.rs b/crates/server/src/service/record.rs index 2c26c0e12..eebd3dfa8 100644 --- a/crates/server/src/service/record.rs +++ b/crates/server/src/service/record.rs @@ -170,6 +170,37 @@ impl RecordService { } } + /// Raw records for a server within the trailing `window`, newest first. + /// The shared read path for alert evaluation: consumers sampling "recent + /// metrics" go through here so a storage change cannot silently detach + /// them from what the recorder writes. + pub async fn query_recent( + db: &DatabaseConnection, + server_id: &str, + window: Duration, + ) -> Result, AppError> { + let since = Utc::now() - window; + Ok(record::Entity::find() + .filter(record::Column::ServerId.eq(server_id)) + .filter(record::Column::Time.gte(since)) + .order_by_desc(record::Column::Time) + .all(db) + .await?) + } + + /// Time of the most recent raw record for a server, if any. + pub async fn latest_record_time( + db: &DatabaseConnection, + server_id: &str, + ) -> Result>, AppError> { + Ok(record::Entity::find() + .filter(record::Column::ServerId.eq(server_id)) + .order_by_desc(record::Column::Time) + .one(db) + .await? + .map(|r| r.time)) + } + /// Query GPU history records for a server within a time range. pub async fn query_gpu_history( db: &DatabaseConnection, diff --git a/crates/server/src/service/rollup.rs b/crates/server/src/service/rollup.rs index 7d2d5dedf..e49e99cdf 100644 --- a/crates/server/src/service/rollup.rs +++ b/crates/server/src/service/rollup.rs @@ -11,6 +11,8 @@ use chrono::{DateTime, Duration, Utc}; +use crate::entity::record; + /// Longest range the raw table serves under `interval = "auto"`; longer /// ranges read the hourly rollup. pub const RAW_WINDOW_MAX_HOURS: i64 = 24; @@ -68,29 +70,123 @@ pub struct MetricColumn { pub name: &'static str, /// How the raw samples fold into the hourly bucket. pub agg: RollupAgg, + /// Alert `rule_type` that reads this column; `None` marks columns that + /// cannot be alerted on (cumulative transfer counters). + pub alert_rule_type: Option<&'static str>, + /// Typed accessor backing the alert read path; nullable columns read 0. + pub read: fn(&record::Model) -> f64, } /// The scalar metric columns, in physical column order. Order matters: the /// generated INSERT lists columns in this order. pub const METRIC_COLUMNS: &[MetricColumn] = &[ - MetricColumn { name: "cpu", agg: RollupAgg::Avg }, - MetricColumn { name: "mem_used", agg: RollupAgg::AvgInt }, - MetricColumn { name: "swap_used", agg: RollupAgg::AvgInt }, - MetricColumn { name: "disk_used", agg: RollupAgg::AvgInt }, - MetricColumn { name: "net_in_speed", agg: RollupAgg::AvgInt }, - MetricColumn { name: "net_out_speed", agg: RollupAgg::AvgInt }, - MetricColumn { name: "net_in_transfer", agg: RollupAgg::MaxInt }, - MetricColumn { name: "net_out_transfer", agg: RollupAgg::MaxInt }, - MetricColumn { name: "load1", agg: RollupAgg::Avg }, - MetricColumn { name: "load5", agg: RollupAgg::Avg }, - MetricColumn { name: "load15", agg: RollupAgg::Avg }, - MetricColumn { name: "tcp_conn", agg: RollupAgg::AvgInt }, - MetricColumn { name: "udp_conn", agg: RollupAgg::AvgInt }, - MetricColumn { name: "process_count", agg: RollupAgg::AvgInt }, - MetricColumn { name: "temperature", agg: RollupAgg::Avg }, - MetricColumn { name: "gpu_usage", agg: RollupAgg::Avg }, + MetricColumn { + name: "cpu", + agg: RollupAgg::Avg, + alert_rule_type: Some("cpu"), + read: |r| r.cpu, + }, + MetricColumn { + name: "mem_used", + agg: RollupAgg::AvgInt, + alert_rule_type: Some("memory"), + read: |r| r.mem_used as f64, + }, + MetricColumn { + name: "swap_used", + agg: RollupAgg::AvgInt, + alert_rule_type: Some("swap"), + read: |r| r.swap_used as f64, + }, + MetricColumn { + name: "disk_used", + agg: RollupAgg::AvgInt, + alert_rule_type: Some("disk"), + read: |r| r.disk_used as f64, + }, + MetricColumn { + name: "net_in_speed", + agg: RollupAgg::AvgInt, + alert_rule_type: Some("net_in_speed"), + read: |r| r.net_in_speed as f64, + }, + MetricColumn { + name: "net_out_speed", + agg: RollupAgg::AvgInt, + alert_rule_type: Some("net_out_speed"), + read: |r| r.net_out_speed as f64, + }, + MetricColumn { + name: "net_in_transfer", + agg: RollupAgg::MaxInt, + alert_rule_type: None, + read: |r| r.net_in_transfer as f64, + }, + MetricColumn { + name: "net_out_transfer", + agg: RollupAgg::MaxInt, + alert_rule_type: None, + read: |r| r.net_out_transfer as f64, + }, + MetricColumn { + name: "load1", + agg: RollupAgg::Avg, + alert_rule_type: Some("load1"), + read: |r| r.load1, + }, + MetricColumn { + name: "load5", + agg: RollupAgg::Avg, + alert_rule_type: Some("load5"), + read: |r| r.load5, + }, + MetricColumn { + name: "load15", + agg: RollupAgg::Avg, + alert_rule_type: Some("load15"), + read: |r| r.load15, + }, + MetricColumn { + name: "tcp_conn", + agg: RollupAgg::AvgInt, + alert_rule_type: Some("tcp_conn"), + read: |r| r.tcp_conn as f64, + }, + MetricColumn { + name: "udp_conn", + agg: RollupAgg::AvgInt, + alert_rule_type: Some("udp_conn"), + read: |r| r.udp_conn as f64, + }, + MetricColumn { + name: "process_count", + agg: RollupAgg::AvgInt, + alert_rule_type: Some("process"), + read: |r| r.process_count as f64, + }, + MetricColumn { + name: "temperature", + agg: RollupAgg::Avg, + alert_rule_type: Some("temperature"), + read: |r| r.temperature.unwrap_or(0.0), + }, + MetricColumn { + name: "gpu_usage", + agg: RollupAgg::Avg, + alert_rule_type: Some("gpu"), + read: |r| r.gpu_usage.unwrap_or(0.0), + }, ]; +/// Metric value an alert rule reads from a raw record. Unknown rule types +/// read 0.0 (the legacy fallback) so a misconfigured rule never fires. +pub fn alert_metric(rec: &record::Model, rule_type: &str) -> f64 { + METRIC_COLUMNS + .iter() + .find(|c| c.alert_rule_type == Some(rule_type)) + .map_or(0.0, |c| (c.read)(rec)) +} + /// Build the hourly rollup upsert. Placeholders: bucket time, window start, /// window end (in that order). pub fn aggregate_hourly_sql() -> String { @@ -204,4 +300,86 @@ mod tests { names.dedup(); assert_eq!(names.len(), METRIC_COLUMNS.len()); } + + fn make_record(cpu: f64, mem_used: i64, load1: f64) -> record::Model { + record::Model { + id: 1, + server_id: "srv-1".to_string(), + time: Utc::now(), + cpu, + mem_used, + swap_used: 0, + disk_used: 0, + net_in_speed: 0, + net_out_speed: 0, + net_in_transfer: 0, + net_out_transfer: 0, + load1, + load5: 0.0, + load15: 0.0, + tcp_conn: 100, + udp_conn: 50, + process_count: 200, + temperature: Some(55.0), + gpu_usage: Some(40.0), + disk_io_json: None, + } + } + + #[test] + fn alert_metric_reads_direct_columns() { + let rec = make_record(85.5, 4_000_000, 1.2); + assert!((alert_metric(&rec, "cpu") - 85.5).abs() < f64::EPSILON); + assert!((alert_metric(&rec, "load1") - 1.2).abs() < f64::EPSILON); + } + + /// Aliased rule types ("memory", "process", "gpu"…) resolve through the + /// descriptor's alert_rule_type, not the column name. + #[test] + fn alert_metric_resolves_rule_type_aliases() { + let rec = make_record(50.0, 8_000_000, 0.0); + assert!((alert_metric(&rec, "memory") - 8_000_000.0).abs() < f64::EPSILON); + assert!((alert_metric(&rec, "process") - 200.0).abs() < f64::EPSILON); + assert!((alert_metric(&rec, "gpu") - 40.0).abs() < f64::EPSILON); + assert!((alert_metric(&rec, "tcp_conn") - 100.0).abs() < f64::EPSILON); + assert!((alert_metric(&rec, "udp_conn") - 50.0).abs() < f64::EPSILON); + assert!((alert_metric(&rec, "temperature") - 55.0).abs() < f64::EPSILON); + } + + /// Unknown rule types and the deliberately non-alertable transfer + /// counters read 0.0 so such rules can never fire. + #[test] + fn alert_metric_unknown_and_non_alertable_read_zero() { + let rec = make_record(99.0, 0, 0.0); + assert!((alert_metric(&rec, "nonexistent") - 0.0).abs() < f64::EPSILON); + assert!((alert_metric(&rec, "net_in_transfer") - 0.0).abs() < f64::EPSILON); + } + + #[test] + fn alert_metric_remaining_variants() { + // Exercise every alertable column not covered by the tests above. + let mut rec = make_record(0.0, 0, 0.0); + rec.swap_used = 1024; + rec.disk_used = 2048; + rec.load5 = 2.5; + rec.load15 = 3.5; + rec.net_in_speed = 111; + rec.net_out_speed = 222; + assert!((alert_metric(&rec, "swap") - 1024.0).abs() < f64::EPSILON); + assert!((alert_metric(&rec, "disk") - 2048.0).abs() < f64::EPSILON); + assert!((alert_metric(&rec, "load5") - 2.5).abs() < f64::EPSILON); + assert!((alert_metric(&rec, "load15") - 3.5).abs() < f64::EPSILON); + assert!((alert_metric(&rec, "net_in_speed") - 111.0).abs() < f64::EPSILON); + assert!((alert_metric(&rec, "net_out_speed") - 222.0).abs() < f64::EPSILON); + } + + #[test] + fn alert_metric_none_temperature_and_gpu_read_zero() { + // None temperature/gpu must coerce to 0.0 via unwrap_or. + let mut rec = make_record(0.0, 0, 0.0); + rec.temperature = None; + rec.gpu_usage = None; + assert!((alert_metric(&rec, "temperature") - 0.0).abs() < f64::EPSILON); + assert!((alert_metric(&rec, "gpu") - 0.0).abs() < f64::EPSILON); + } } From e3e868da6d98617fc621ba19598f0fa49e1199f2 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 23:28:59 +0800 Subject: [PATCH 31/44] refactor(server): normalize public metrics mapping through one row shape get_server_metrics hand-copied the same 16-field PublicMetricsPoint mapping for the raw and hourly query branches. Make the identical- column-set fact code: record_hourly::Model converts into record::Model, QueryHistoryResult::into_rows() erases the resolution for consumers that don't care which table served them, and the public mapping shrinks to a single field list. --- crates/server/src/entity/record_hourly.rs | 30 +++++++++ crates/server/src/service/public_status.rs | 72 ++++++++-------------- crates/server/src/service/record.rs | 12 ++++ 3 files changed, 67 insertions(+), 47 deletions(-) diff --git a/crates/server/src/entity/record_hourly.rs b/crates/server/src/entity/record_hourly.rs index 1f232dd11..9460a1f4e 100644 --- a/crates/server/src/entity/record_hourly.rs +++ b/crates/server/src/entity/record_hourly.rs @@ -33,3 +33,33 @@ pub struct Model { pub enum Relation {} impl ActiveModelBehavior for ActiveModel {} + +/// `records_hourly` deliberately mirrors the `records` column set (the rollup +/// policy manages both), so an hourly row converts losslessly into the raw-row +/// shape for consumers that don't care which resolution served them. +impl From for super::record::Model { + fn from(m: Model) -> Self { + super::record::Model { + id: m.id, + server_id: m.server_id, + time: m.time, + cpu: m.cpu, + mem_used: m.mem_used, + swap_used: m.swap_used, + disk_used: m.disk_used, + net_in_speed: m.net_in_speed, + net_out_speed: m.net_out_speed, + net_in_transfer: m.net_in_transfer, + net_out_transfer: m.net_out_transfer, + load1: m.load1, + load5: m.load5, + load15: m.load15, + tcp_conn: m.tcp_conn, + udp_conn: m.udp_conn, + process_count: m.process_count, + temperature: m.temperature, + gpu_usage: m.gpu_usage, + disk_io_json: m.disk_io_json, + } + } +} diff --git a/crates/server/src/service/public_status.rs b/crates/server/src/service/public_status.rs index 263a8a380..10428649e 100644 --- a/crates/server/src/service/public_status.rs +++ b/crates/server/src/service/public_status.rs @@ -21,7 +21,7 @@ use crate::entity::{ use crate::error::AppError; use crate::service::agent_manager::aggregate_disk_io; use crate::service::ip_quality::IpQualityService; -use crate::service::record::{QueryHistoryResult, RecordService}; +use crate::service::record::RecordService; use crate::service::uptime::{UptimeDailyEntry, UptimeService}; // --------------------------------------------------------------------------- @@ -641,52 +641,30 @@ pub async fn get_server_metrics( let from = clamp_public_metrics_from(range.from, range.to, &range.interval)?; let result = RecordService::query_history(db, id, from, range.to, &range.interval).await?; - let points = match result { - QueryHistoryResult::Raw(records) => records - .into_iter() - .map(|r| PublicMetricsPoint { - time: r.time.to_rfc3339(), - cpu: r.cpu, - mem_used: r.mem_used, - disk_used: r.disk_used, - net_in_speed: r.net_in_speed, - net_out_speed: r.net_out_speed, - net_in_transfer: r.net_in_transfer, - net_out_transfer: r.net_out_transfer, - load1: r.load1, - load5: r.load5, - load15: r.load15, - tcp_conn: r.tcp_conn, - udp_conn: r.udp_conn, - process_count: r.process_count, - temperature: r.temperature, - gpu_usage: r.gpu_usage, - }) - .collect(), - QueryHistoryResult::Hourly(records) => records - .into_iter() - .map(|r| PublicMetricsPoint { - time: r.time.to_rfc3339(), - cpu: r.cpu, - mem_used: r.mem_used, - disk_used: r.disk_used, - net_in_speed: r.net_in_speed, - net_out_speed: r.net_out_speed, - net_in_transfer: r.net_in_transfer, - net_out_transfer: r.net_out_transfer, - load1: r.load1, - load5: r.load5, - load15: r.load15, - tcp_conn: r.tcp_conn, - udp_conn: r.udp_conn, - process_count: r.process_count, - temperature: r.temperature, - gpu_usage: r.gpu_usage, - }) - .collect(), - }; - - Ok(points) + // Both resolutions share the raw-row shape, so one field mapping serves + // raw and hourly queries alike. + Ok(result + .into_rows() + .into_iter() + .map(|r| PublicMetricsPoint { + time: r.time.to_rfc3339(), + cpu: r.cpu, + mem_used: r.mem_used, + disk_used: r.disk_used, + net_in_speed: r.net_in_speed, + net_out_speed: r.net_out_speed, + net_in_transfer: r.net_in_transfer, + net_out_transfer: r.net_out_transfer, + load1: r.load1, + load5: r.load5, + load15: r.load15, + tcp_conn: r.tcp_conn, + udp_conn: r.udp_conn, + process_count: r.process_count, + temperature: r.temperature, + gpu_usage: r.gpu_usage, + }) + .collect()) } /// 90-day uptime band for a scoped server. diff --git a/crates/server/src/service/record.rs b/crates/server/src/service/record.rs index eebd3dfa8..1c806c3ea 100644 --- a/crates/server/src/service/record.rs +++ b/crates/server/src/service/record.rs @@ -334,6 +334,18 @@ pub enum QueryHistoryResult { Hourly(Vec), } +impl QueryHistoryResult { + /// Rows in the raw-row shape regardless of resolution (the two tables + /// share one column set). For consumers that don't care which table + /// served the query. + pub fn into_rows(self) -> Vec { + match self { + QueryHistoryResult::Raw(rows) => rows, + QueryHistoryResult::Hourly(rows) => rows.into_iter().map(Into::into).collect(), + } + } +} + #[cfg(test)] mod tests { use super::*; From fcf7fcddb003bc4fdd19daf4ec77cd0f30b100ea Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 23:35:51 +0800 Subject: [PATCH 32/44] refactor(web): render metric charts from a display descriptor list MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Nine hand-written MetricsChart JSX blocks each re-stated a metric's dataKey, color, unit, domain, byte formatting and availability gate. Declare them once in METRIC_CHART_SPECS and render through one loop, so adding a chart is a descriptor row plus its i18n keys. RealtimeDataPoint becomes an alias of MetricChartRow — the WS feed and the records API already converged on one chart row shape; now the types say so. --- .../status/server-detail-content.tsx | 120 +++--------------- apps/web/src/hooks/use-realtime-metrics.ts | 18 +-- apps/web/src/lib/metric-chart-model.ts | 86 +++++++++++++ 3 files changed, 110 insertions(+), 114 deletions(-) diff --git a/apps/web/src/components/status/server-detail-content.tsx b/apps/web/src/components/status/server-detail-content.tsx index 677336f7e..c156547dc 100644 --- a/apps/web/src/components/status/server-detail-content.tsx +++ b/apps/web/src/components/status/server-detail-content.tsx @@ -29,6 +29,8 @@ import { buildGpuChartRows, deriveNetworkLabels, type GpuRecordAggregated, + METRIC_CHART_SPECS, + type MetricChartSpec, makeTickFormatter, makeTooltipFormatter, toMetricChartRow, @@ -401,6 +403,13 @@ function MetricsTabContent({ const { t } = useTranslation('servers') const hasGpuTemp = gpuChartData.some((d) => 'gpu_temp' in d && d.gpu_temp != null) const isPublic = variant === 'public' + const gates: Record, boolean> = { + gpu: availableMetrics.gpu, + // GPU temp series is admin-only; the public surface does not expose it, + // so the gate also requires a non-empty data key. + gpuTemp: availableMetrics.gpu && hasGpuTemp, + temperature: availableMetrics.temperature + } return ( <> @@ -419,111 +428,22 @@ function MetricsTabContent({
- - - - formatBytes(v)} - formatTime={formatTime} - formatTooltipLabel={formatTooltipLabel} - formatValue={(v) => formatBytes(v)} - title={t('chart_net_in')} - xAxisInterval={xAxisInterval} - /> - formatBytes(v)} - formatTime={formatTime} - formatTooltipLabel={formatTooltipLabel} - formatValue={(v) => formatBytes(v)} - title={t('chart_net_out')} - xAxisInterval={xAxisInterval} - /> - - - {availableMetrics.temperature && ( - - )} - - {availableMetrics.gpu && ( + {METRIC_CHART_SPECS.filter((spec) => !spec.gate || gates[spec.gate]).map((spec) => ( - )} - {/* GPU temp series is admin-only; the public surface does not - expose it, so we gate the chart on a non-empty data key. */} - {availableMetrics.gpu && hasGpuTemp && ( - - )} + ))}
{availableMetrics.diskIo && ( diff --git a/apps/web/src/hooks/use-realtime-metrics.ts b/apps/web/src/hooks/use-realtime-metrics.ts index ec1853e0e..305e7d98a 100644 --- a/apps/web/src/hooks/use-realtime-metrics.ts +++ b/apps/web/src/hooks/use-realtime-metrics.ts @@ -1,7 +1,7 @@ import type { QueryClient } from '@tanstack/react-query' import { useQueryClient } from '@tanstack/react-query' import { useEffect, useRef, useState } from 'react' -import { pct } from '@/lib/metric-chart-model' +import { type MetricChartRow, pct } from '@/lib/metric-chart-model' import { readLiveServers, type ServerMetrics, subscribeLiveServers } from '@/lib/server-catalog' const MAX_BUFFER_SIZE = 200 @@ -31,19 +31,9 @@ function getQueryClientBufferCache( return cache } -export interface RealtimeDataPoint { - cpu: number - disk_pct: number - load1: number - load5: number - load15: number - memory_pct: number - net_in_speed: number - net_in_transfer: number - net_out_speed: number - net_out_transfer: number - timestamp: string -} +/** Realtime points are ordinary chart rows: the WS feed and the records API + * converge on one row shape so every chart reads the same keys. */ +export type RealtimeDataPoint = MetricChartRow export function toRealtimeDataPoint(metrics: ServerMetrics): RealtimeDataPoint { return { diff --git a/apps/web/src/lib/metric-chart-model.ts b/apps/web/src/lib/metric-chart-model.ts index 16eb1fc97..fc30c4767 100644 --- a/apps/web/src/lib/metric-chart-model.ts +++ b/apps/web/src/lib/metric-chart-model.ts @@ -41,6 +41,92 @@ export type MetricChartRow = { timestamp: string } +/** + * Display spec for one metric chart. Everything the Metrics tab needs to + * render a series — data key, label, color, unit, domain, byte formatting, + * data source and availability gate — lives in this one table, so adding a + * chart means adding a row here instead of a hand-written JSX block. + */ +export interface MetricChartSpec { + /** Format axis ticks and tooltip values with formatBytes. */ + bytes?: boolean + /** CSS color value for the area series. */ + color: string + /** Row key in the chart data. */ + dataKey: string + domain?: [number, number] + /** Availability gate; ungated charts always render. */ + gate?: 'gpu' | 'gpuTemp' | 'temperature' + /** i18n key in the `servers` namespace. */ + labelKey: string + /** Which row array feeds the chart. */ + source: 'gpu' | 'metrics' + unit?: string +} + +export const METRIC_CHART_SPECS: MetricChartSpec[] = [ + { + color: 'var(--color-chart-1)', + dataKey: 'cpu', + domain: [0, 100], + labelKey: 'chart_cpu', + source: 'metrics', + unit: '%' + }, + { + color: 'var(--color-chart-2)', + dataKey: 'memory_pct', + domain: [0, 100], + labelKey: 'chart_memory', + source: 'metrics', + unit: '%' + }, + { + color: 'var(--color-chart-3)', + dataKey: 'disk_pct', + domain: [0, 100], + labelKey: 'chart_disk', + source: 'metrics', + unit: '%' + }, + { bytes: true, color: 'var(--color-chart-4)', dataKey: 'net_in_speed', labelKey: 'chart_net_in', source: 'metrics' }, + { + bytes: true, + color: 'var(--color-chart-5)', + dataKey: 'net_out_speed', + labelKey: 'chart_net_out', + source: 'metrics' + }, + { color: 'var(--color-chart-1)', dataKey: 'load1', labelKey: 'chart_load', source: 'metrics' }, + { + color: 'var(--color-chart-4)', + dataKey: 'temperature', + gate: 'temperature', + labelKey: 'chart_temperature', + source: 'metrics', + unit: '°C' + }, + { + color: 'var(--color-chart-5)', + dataKey: 'gpu_usage', + domain: [0, 100], + gate: 'gpu', + labelKey: 'chart_gpu_usage', + source: 'gpu', + unit: '%' + }, + // GPU temp is admin-only: the public surface never populates gpu_temp, so + // the gpuTemp gate (non-empty data key) keeps the chart off that surface. + { + color: 'var(--color-chart-2)', + dataKey: 'gpu_temp', + gate: 'gpuTemp', + labelKey: 'chart_gpu_temp', + source: 'gpu', + unit: '°C' + } +] + /** The one percentage guard: a missing/zero total renders as 0%, never NaN/Infinity. */ export function pct(used: number, total: number): number { return total > 0 ? (used / total) * 100 : 0 From 61e841e70e159db95fe375dbe83f47ffa8fbbf0b Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sat, 11 Jul 2026 23:37:06 +0800 Subject: [PATCH 33/44] docs: record the rollup descriptor decision and glossary terms ADR-0003 captures why scalar metric columns are declared once in service::rollup (generated SQL pinned by a golden test, alert reads via the shared descriptor, resolution-erasing row conversion) and why full struct codegen and OpenAPI-derived chart specs were rejected. CONTEXT.md gains the Rollup policy and Metric column descriptor terms. --- CONTEXT.md | 2 + ...tric-columns-owned-by-rollup-descriptor.md | 65 +++++++++++++++++++ 2 files changed, 67 insertions(+) create mode 100644 docs/adr/0003-metric-columns-owned-by-rollup-descriptor.md diff --git a/CONTEXT.md b/CONTEXT.md index 4f8b5ec05..cc2e792e1 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -6,3 +6,5 @@ - **Standalone public network page** — `/status/network/$serverId`; exists only as a fallback for status pages configured with `show_server_detail=false` but `show_network=true` (see ADR-0001). - **Live metrics (实时指标帧)** — the partial projection of a server carried by WS `update` frames: the fields an agent report can populate (usage, speeds, loads, connection counts), plus `name` kept purely for decoder compatibility. Static facts are absent from the wire; clients keep their cached values when merging (see ADR-0002). _Avoid_: partial status, update payload. - **Full server status** — the complete ~35-field snapshot (`ServerStatus`) carried by `full_sync` and REST. The only sources allowed to seed the catalog or overwrite static facts (totals, os, tags, geo, enrollment). +- **Rollup policy (降采样策略)** — how raw metric records compress over time: per-column aggregation (AVG vs window-end MAX), the raw/hourly table switch point, and which columns alerts may read. Owned by `service::rollup`; declared per column in the **metric column descriptor** (see ADR-0003). _Avoid_: aggregation logic, downsampling code. +- **Metric column descriptor** — the one-row-per-column table (`METRIC_COLUMNS`) declaring a scalar metric's SQL name, rollup aggregation, alert rule type, and typed accessor. Adding a scalar metric means adding one descriptor row (plus entity/migration, which the compiler enforces). The web analogue is `METRIC_CHART_SPECS` for chart display. diff --git a/docs/adr/0003-metric-columns-owned-by-rollup-descriptor.md b/docs/adr/0003-metric-columns-owned-by-rollup-descriptor.md new file mode 100644 index 000000000..0433afad7 --- /dev/null +++ b/docs/adr/0003-metric-columns-owned-by-rollup-descriptor.md @@ -0,0 +1,65 @@ +# ADR-0003: Scalar metric columns are owned by the rollup descriptor + +## Status + +Accepted (2026-07-11) + +## Context + +A scalar metric's knowledge was re-encoded in several places that nothing +cross-checked: the hourly rollup SQL string listed every column two to three +times and was the only place the AVG-vs-MAX aggregation decision lived; +`query_history` hardcoded the 24h raw/hourly switch; alert evaluation kept a +private read path (`check_threshold` querying the records table directly) plus +a second rule-type→column `match`. Adding one metric column meant hand-editing +~16 sites across 9 files, and the two easiest sites to miss (the SQL string, +the public DTO maps) had no compiler protection. + +## Decision + +`service::rollup` owns the rollup policy. Every scalar column of +`records`/`records_hourly` is declared once in `METRIC_COLUMNS` with its SQL +name, `RollupAgg` (AVG / CAST-AVG / window-end MAX), optional alert rule type, +and a typed accessor. From that table we derive: + +- the hourly rollup upsert (`aggregate_hourly_sql()`), pinned byte-for-byte to + the original hand-written SQL by a golden test; +- the alert metric read (`alert_metric`), with transfer counters deliberately + non-alertable (`alert_rule_type: None`); +- the raw/hourly switch (`select_history_table`, `RAW_WINDOW_MAX_HOURS`). + +Alert evaluation reads through `RecordService::query_recent` / +`latest_record_time` instead of querying tables directly, so a storage change +cannot silently detach alerting from what the recorder writes. Because +`records_hourly` mirrors the `records` column set, `record_hourly::Model` +converts into `record::Model` and `QueryHistoryResult::into_rows()` erases the +resolution for consumers that don't care (the public metrics DTO maps once, +not per branch). + +On the web, chart display specs (dataKey, color, unit, domain, byte +formatting, availability gate) live in `METRIC_CHART_SPECS` and render through +one loop. + +## Alternatives considered + +- **Proc-macro / build-script codegen of the structs themselves** (SystemReport, + entities, DTOs from one schema): kills the remaining duplication but adds a + compile-time machinery cost far above the ~yearly rate of adding metrics. + The struct field lists are already compiler-checked; only the SQL and + string-keyed maps were not. +- **Deriving the web metric list from the OpenAPI schema**: the generated + `api-types.ts` already mirrors field names/types; what the charts need is + display knowledge (color, unit, gate) that OpenAPI does not carry, so a + hand-written display descriptor is the honest owner. + +## Consequences + +- Adding a scalar metric: one `METRIC_COLUMNS` row + entity/migration + + protocol/DTO fields (compiler-enforced) + one `METRIC_CHART_SPECS` row and + i18n keys for a chart. The golden SQL test must be updated deliberately, + which is the point — aggregation choices are reviewed, not implied. +- `disk_io_json` stays the documented exception (per-device JSON folded in + Rust after the SQL pass); GPU detail keeps its own sub-table with no rollup. +- The descriptor's accessor field ties it to `record::Model`; if raw storage + ever splits from the entity shape, the descriptor is the single seam to + re-point. From 90bbbbc06e7c412b44f21f8849708a5922d8a46b Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sun, 12 Jul 2026 00:23:49 +0800 Subject: [PATCH 34/44] refactor(agent): extract the Docker availability FSM into a subsystem The retry/demote state machine (4 fields, 5 methods) lived inside the shared ConnectionRuntime, and its transitions wrote FeaturesUpdate / DockerUnavailable straight to the WebSocket sink, so exercising them required a live daemon plus a recording sink. DockerSubsystem owns detection, promote/demote, stats-poll gating and the unavailable replies; methods return the message to send instead of writing it, and the availability transitions now have direct unit tests (announce-once demotion, msg_id correlation, stats-interval lifecycle). The dispatcher match collapses to one delegating arm; dispatch-level harness coverage in runtime::tests is unchanged. Deliberately not a full Subsystem trait across all message families: ping/terminal/file/firewall already sit behind their own managers, and a uniform trait over heterogeneous shapes would be a shallow interface serving a single hypothetical seam. --- crates/agent/src/reporter/docker_subsystem.rs | 390 ++++++++++++++++++ crates/agent/src/reporter/mod.rs | 19 +- crates/agent/src/reporter/runtime.rs | 283 ++----------- 3 files changed, 425 insertions(+), 267 deletions(-) create mode 100644 crates/agent/src/reporter/docker_subsystem.rs diff --git a/crates/agent/src/reporter/docker_subsystem.rs b/crates/agent/src/reporter/docker_subsystem.rs new file mode 100644 index 000000000..c586879c7 --- /dev/null +++ b/crates/agent/src/reporter/docker_subsystem.rs @@ -0,0 +1,390 @@ +//! Docker availability state machine, connection-scoped. +//! +//! One WebSocket connection owns one [`DockerSubsystem`]: daemon detection, +//! the retry/demote lifecycle, stats-poll gating, and the unavailable replies +//! for Docker requests all live here. The shared runtime routes Docker frames +//! in and forwards the returned [`AgentMessage`]s out — nothing more. Methods +//! return the message to send instead of writing to the WebSocket sink, so +//! every availability transition is unit-testable without a Docker daemon. + +use std::sync::Arc; +use std::time::Duration; + +use serverbee_common::protocol::{AgentMessage, ServerMessage}; +use tokio::sync::mpsc; + +use crate::capability_grants::CapabilityAuthority; +use crate::docker::DockerManager; + +const DOCKER_RETRY_SECS: u64 = 30; + +/// Which housekeeping event [`DockerSubsystem::tick`] resolved to: poll +/// container stats, or retry the daemon connection. +pub(super) enum DockerTick { + PollStats, + Retry, +} + +pub(super) struct DockerSubsystem { + capabilities: Arc, + /// Sender cloned into every spawned `DockerManager` task; the matching + /// receiver is drained by the select! loop in `mod.rs`. + tx: mpsc::Sender, + manager: Option, + available: bool, + stats_interval: Option, + retry_interval: tokio::time::Interval, +} + +impl DockerSubsystem { + pub(super) fn new( + tx: mpsc::Sender, + capabilities: Arc, + ) -> Self { + // First retry fires one full period out — connection start already + // probes explicitly, so there is no immediate tick to consume. + let retry_period = Duration::from_secs(DOCKER_RETRY_SECS); + let retry_interval = + tokio::time::interval_at(tokio::time::Instant::now() + retry_period, retry_period); + Self { + capabilities, + tx, + manager: None, + available: false, + stats_interval: None, + retry_interval, + } + } + + /// Feature list advertised in `SystemInfo`, derived from what is live. + pub(super) fn features(&self) -> Vec { + if self.available { + vec!["docker".to_string()] + } else { + Vec::new() + } + } + + /// Initial daemon detection at connection start. Absence is informational + /// — the retry tick keeps looking. + pub(super) async fn probe(&mut self) { + match DockerManager::try_new(self.tx.clone(), Arc::clone(&self.capabilities)) { + Ok(dm) => match dm.verify_connection().await { + Ok(()) => { + tracing::info!("Docker daemon connected"); + self.available = true; + self.manager = Some(dm); + } + Err(e) => { + tracing::info!("Docker daemon not available: {e}"); + } + }, + Err(e) => { + tracing::info!("Docker not available: {e}"); + } + } + } + + /// Wait for the next housekeeping event: a stats-poll tick while the + /// daemon is connected, or a retry tick while it is absent. Pending when + /// connected with stats polling off. + pub(super) async fn tick(&mut self) -> DockerTick { + if self.manager.is_some() { + match self.stats_interval.as_mut() { + Some(iv) => { + iv.tick().await; + DockerTick::PollStats + } + None => std::future::pending().await, + } + } else { + self.retry_interval.tick().await; + DockerTick::Retry + } + } + + /// Poll container stats once; a failing daemon demotes the subsystem. + /// Returns the `FeaturesUpdate` to send when availability changed. + pub(super) async fn poll_stats(&mut self) -> Option { + if let Some(dm) = self.manager.as_mut() + && let Err(e) = dm.poll_stats().await + { + tracing::warn!("Docker stats polling failed: {e}"); + return self.demote(); + } + None + } + + /// Retry the daemon connection; success returns the `FeaturesUpdate` + /// re-advertising the feature. + pub(super) async fn retry(&mut self) -> Option { + tracing::debug!("Retrying Docker connection..."); + match DockerManager::try_new(self.tx.clone(), Arc::clone(&self.capabilities)) { + Ok(dm) => match dm.verify_connection().await { + Ok(()) => { + tracing::info!("Docker daemon now available"); + self.manager = Some(dm); + self.available = true; + return Some(AgentMessage::FeaturesUpdate { + features: vec!["docker".to_string()], + }); + } + Err(e) => { + tracing::debug!("Docker still not available: {e}"); + } + }, + Err(e) => { + tracing::debug!("Docker still not available: {e}"); + } + } + None + } + + /// Handle one Docker-family server message (stats control or a request + /// forwarded to the daemon). Returns the reply to send, if any. + pub(super) async fn handle(&mut self, msg: ServerMessage) -> Option { + match msg { + ServerMessage::DockerStartStats { interval_secs } => { + if self.manager.is_some() { + let secs = interval_secs.max(1); + tracing::info!("Starting Docker stats polling every {secs}s"); + let mut iv = tokio::time::interval(Duration::from_secs(secs as u64)); + iv.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + self.stats_interval = Some(iv); + None + } else { + tracing::warn!("DockerStartStats received but Docker is not available"); + Some(AgentMessage::DockerUnavailable { msg_id: None }) + } + } + ServerMessage::DockerStopStats => { + tracing::info!("Stopping Docker stats polling"); + self.stats_interval = None; + None + } + msg => { + if let Some(dm) = self.manager.as_mut() { + if let Err(e) = dm.handle_server_message(msg).await { + tracing::warn!("Docker runtime became unavailable: {e}"); + return self.demote(); + } + None + } else { + tracing::warn!("Docker message received but Docker is not available"); + Some(AgentMessage::DockerUnavailable { + msg_id: docker_request_msg_id(&msg), + }) + } + } + } + } + + /// Drop the daemon connection and stop stats polling. Returns the empty + /// `FeaturesUpdate` when the feature was previously advertised. + fn demote(&mut self) -> Option { + if let Some(dm) = self.manager.as_mut() { + dm.cleanup(); + } + self.manager = None; + self.stats_interval = None; + + if self.available { + self.available = false; + Some(AgentMessage::FeaturesUpdate { features: vec![] }) + } else { + None + } + } + + /// Whether a stats-poll interval is armed — observability for the + /// dispatch harness in `runtime::tests`. + #[cfg(test)] + pub(super) fn stats_polling_active(&self) -> bool { + self.stats_interval.is_some() + } + + /// Pre-arm stats polling so tests can observe it being cleared. + #[cfg(test)] + pub(super) fn arm_stats_polling_for_test(&mut self) { + self.stats_interval = Some(tokio::time::interval(Duration::from_secs(60))); + } + + /// Tear down the daemon connection without emitting anything; part of + /// connection shutdown where the socket is already gone. + pub(super) fn cleanup(&mut self) { + if let Some(dm) = self.manager.as_mut() { + dm.cleanup(); + } + } +} + +/// The `msg_id` a `DockerUnavailable` reply should carry so the server can +/// correlate it with the request; log/event stream variants have none. +fn docker_request_msg_id(msg: &ServerMessage) -> Option { + match msg { + ServerMessage::DockerListContainers { msg_id } + | ServerMessage::DockerContainerAction { msg_id, .. } + | ServerMessage::DockerGetInfo { msg_id } + | ServerMessage::DockerListNetworks { msg_id } + | ServerMessage::DockerListVolumes { msg_id } => Some(msg_id.clone()), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn make_subsystem() -> DockerSubsystem { + let (tx, _rx) = mpsc::channel(8); + let dir = tempfile::tempdir().expect("tempdir"); + let capabilities = CapabilityAuthority::new( + serverbee_common::constants::CAP_DEFAULT, + dir.path().join("grants.json"), + ); + DockerSubsystem::new(tx, capabilities) + } + + #[test] + fn test_docker_request_msg_id_extracts_only_request_variants() { + // Variants that carry a msg_id return it... + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerListContainers { + msg_id: "a".to_string() + }), + Some("a".to_string()) + ); + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerGetInfo { + msg_id: "b".to_string() + }), + Some("b".to_string()) + ); + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerListNetworks { + msg_id: "c".to_string() + }), + Some("c".to_string()) + ); + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerListVolumes { + msg_id: "d".to_string() + }), + Some("d".to_string()) + ); + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerContainerAction { + msg_id: "e".to_string(), + container_id: "cid".to_string(), + action: serverbee_common::docker_types::DockerAction::Restart { timeout: None }, + }), + Some("e".to_string()) + ); + // ...non-request docker variants return None. + assert_eq!(docker_request_msg_id(&ServerMessage::DockerStopStats), None); + assert_eq!(docker_request_msg_id(&ServerMessage::Ping), None); + } + + #[test] + fn test_docker_request_msg_id_none_for_log_and_event_variants() { + // Streaming control variants carry no request msg_id -> None. + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerLogsStart { + session_id: "s".to_string(), + container_id: "c".to_string(), + tail: None, + follow: false, + }), + None + ); + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerLogsStop { + session_id: "s".to_string(), + }), + None + ); + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerEventsStart), + None + ); + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerEventsStop), + None + ); + assert_eq!( + docker_request_msg_id(&ServerMessage::DockerStartStats { interval_secs: 5 }), + None + ); + } + + #[tokio::test] + async fn test_start_stats_unavailable_replies_unavailable_without_arming_interval() { + let mut docker = make_subsystem(); + let reply = docker + .handle(ServerMessage::DockerStartStats { interval_secs: 5 }) + .await; + assert!(matches!( + reply, + Some(AgentMessage::DockerUnavailable { msg_id: None }) + )); + assert!(docker.stats_interval.is_none()); + } + + #[tokio::test] + async fn test_stop_stats_clears_interval_silently() { + let mut docker = make_subsystem(); + docker.stats_interval = Some(tokio::time::interval(Duration::from_secs(1))); + let reply = docker.handle(ServerMessage::DockerStopStats).await; + assert!(reply.is_none()); + assert!(docker.stats_interval.is_none()); + } + + #[tokio::test] + async fn test_request_unavailable_reply_carries_msg_id() { + let mut docker = make_subsystem(); + let reply = docker + .handle(ServerMessage::DockerGetInfo { + msg_id: "req-7".into(), + }) + .await; + match reply { + Some(AgentMessage::DockerUnavailable { msg_id }) => { + assert_eq!(msg_id.as_deref(), Some("req-7")); + } + other => panic!("expected DockerUnavailable, got {other:?}"), + } + } + + /// Demotion is announce-once: the empty FeaturesUpdate goes out only when + /// the feature was previously advertised, and a second demotion is silent. + #[tokio::test] + async fn test_demote_emits_features_update_only_when_previously_available() { + let mut docker = make_subsystem(); + assert!( + docker.demote().is_none(), + "never-available demote is silent" + ); + + docker.available = true; + docker.stats_interval = Some(tokio::time::interval(Duration::from_secs(1))); + match docker.demote() { + Some(AgentMessage::FeaturesUpdate { features }) => { + assert!(features.is_empty(), "demotion clears the feature list"); + } + other => panic!("expected FeaturesUpdate, got {other:?}"), + } + assert!( + docker.stats_interval.is_none(), + "demotion stops stats polling" + ); + assert!(docker.demote().is_none(), "second demote is silent"); + } + + #[tokio::test] + async fn test_features_reflect_availability() { + let mut docker = make_subsystem(); + assert!(docker.features().is_empty()); + docker.available = true; + assert_eq!(docker.features(), vec!["docker".to_string()]); + } +} diff --git a/crates/agent/src/reporter/mod.rs b/crates/agent/src/reporter/mod.rs index c6676eb07..815e02537 100644 --- a/crates/agent/src/reporter/mod.rs +++ b/crates/agent/src/reporter/mod.rs @@ -1,3 +1,4 @@ +mod docker_subsystem; mod file_ops; mod runtime; mod wire; @@ -16,7 +17,8 @@ use tokio::time::interval; use tokio_tungstenite::connect_async; use tokio_tungstenite::tungstenite::Message; -use runtime::{ConnectionRuntime, DockerTick}; +use docker_subsystem::DockerTick; +use runtime::ConnectionRuntime; use wire::send_msg; use crate::capability_grants::CapabilityAuthority; @@ -188,8 +190,8 @@ impl Reporter { Arc::clone(&capabilities), Arc::clone(&self.firewall_manager), ); - runtime.probe_docker().await; - let features = runtime.features(); + runtime.docker.probe().await; + let features = runtime.docker.features(); // Send SystemInfo let mut collector = Collector::new( @@ -419,10 +421,13 @@ impl Reporter { } // Docker housekeeping: stats polling while the daemon is // connected, reconnect retry while it is absent. - tick = runtime.docker_tick() => { - match tick { - DockerTick::PollStats => runtime.poll_docker_stats(&mut write).await?, - DockerTick::Retry => runtime.retry_docker(&mut write).await?, + tick = runtime.docker.tick() => { + let reply = match tick { + DockerTick::PollStats => runtime.docker.poll_stats().await, + DockerTick::Retry => runtime.docker.retry().await, + }; + if let Some(msg) = reply { + send_msg(&mut write, &msg).await?; } } // IP change detection — spawned off the WS hot path so a diff --git a/crates/agent/src/reporter/runtime.rs b/crates/agent/src/reporter/runtime.rs index 9b19db0d1..c029eb754 100644 --- a/crates/agent/src/reporter/runtime.rs +++ b/crates/agent/src/reporter/runtime.rs @@ -1,11 +1,12 @@ //! Connection-scoped command runtime. //! -//! One WebSocket connection owns one [`ConnectionRuntime`]: every manager, -//! outbound channel sender, and piece of Docker lifecycle state that executing -//! a `ServerMessage` can touch lives here, so the dispatcher is a method on -//! the state it drives instead of a free function threading a dozen borrows. -//! `mod.rs` keeps transport concerns only (connect/backoff, the select! loop, -//! report and IP timers); file replies stay in `file_ops`, framing in `wire`. +//! One WebSocket connection owns one [`ConnectionRuntime`]: every manager +//! and outbound channel sender that executing a `ServerMessage` can touch +//! lives here, so the dispatcher is a method on the state it drives instead +//! of a free function threading a dozen borrows. `mod.rs` keeps transport +//! concerns only (connect/backoff, the select! loop, report and IP timers); +//! file replies stay in `file_ops`, framing in `wire`, and the Docker +//! availability state machine in `docker_subsystem`. use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; @@ -18,11 +19,11 @@ use serverbee_common::types::{NetworkProbeResultData, PingResult}; use tokio::sync::mpsc; use tokio_tungstenite::tungstenite::Message; +use super::docker_subsystem::DockerSubsystem; use super::wire::send_msg; use super::{emit_upgrade_failure, file_ops, perform_upgrade}; use crate::capability_grants::CapabilityAuthority; use crate::config::{AgentConfig, UpgradeConfig}; -use crate::docker::DockerManager; use crate::file_manager::{FileEvent, FileManager}; use crate::firewall::FirewallManager; use crate::ip_quality::{RunResult, UnlockChecker}; @@ -30,19 +31,10 @@ use crate::network_prober::NetworkProber; use crate::pinger::PingManager; use crate::terminal::{TerminalEvent, TerminalManager}; -const DOCKER_RETRY_SECS: u64 = 30; - /// Process-wide single-flight latch for binary self-upgrade. A duplicate /// Upgrade command must be rejected without disturbing the running one. static UPGRADE_IN_PROGRESS: AtomicBool = AtomicBool::new(false); -/// Which Docker housekeeping event [`ConnectionRuntime::docker_tick`] resolved -/// to: poll container stats, or retry the daemon connection. -pub(super) enum DockerTick { - PollStats, - Retry, -} - /// Receiver halves for every event stream the runtime's managers emit. The /// select! loop in `mod.rs` drains these; the matching senders live inside /// [`ConnectionRuntime`] so a handler can never outlive its channel. @@ -66,11 +58,7 @@ pub(super) struct ConnectionRuntime { pub(super) unlock_checker: UnlockChecker, pub(super) cmd_result_tx: mpsc::Sender, file_tx: mpsc::Sender, - docker_tx: mpsc::Sender, - docker_manager: Option, - docker_available: bool, - docker_stats_interval: Option, - docker_retry_interval: tokio::time::Interval, + pub(super) docker: DockerSubsystem, upgrade_cfg: UpgradeConfig, firewall_manager: Arc, } @@ -100,15 +88,10 @@ impl ConnectionRuntime { let file_manager = FileManager::new(config.file.clone(), Arc::clone(&capabilities)); let (docker_tx, docker_rx) = mpsc::channel::(256); + let docker = DockerSubsystem::new(docker_tx, Arc::clone(&capabilities)); let (cmd_result_tx, cmd_result_rx) = mpsc::channel::(32); - // First retry fires one full period out — connection start already - // probes explicitly, so there is no immediate tick to consume. - let retry_period = Duration::from_secs(DOCKER_RETRY_SECS); - let docker_retry_interval = - tokio::time::interval_at(tokio::time::Instant::now() + retry_period, retry_period); - let runtime = Self { capabilities, ping_manager, @@ -118,11 +101,7 @@ impl ConnectionRuntime { unlock_checker, cmd_result_tx, file_tx, - docker_tx, - docker_manager: None, - docker_available: false, - docker_stats_interval: None, - docker_retry_interval, + docker, upgrade_cfg: config.upgrade.clone(), firewall_manager, }; @@ -138,114 +117,6 @@ impl ConnectionRuntime { (runtime, events) } - /// Feature list advertised in `SystemInfo`, derived from what is live. - pub(super) fn features(&self) -> Vec { - if self.docker_available { - vec!["docker".to_string()] - } else { - Vec::new() - } - } - - /// Initial Docker detection at connection start. Absence is informational - /// — the retry tick keeps looking. - pub(super) async fn probe_docker(&mut self) { - match DockerManager::try_new(self.docker_tx.clone(), Arc::clone(&self.capabilities)) { - Ok(dm) => match dm.verify_connection().await { - Ok(()) => { - tracing::info!("Docker daemon connected"); - self.docker_available = true; - self.docker_manager = Some(dm); - } - Err(e) => { - tracing::info!("Docker daemon not available: {e}"); - } - }, - Err(e) => { - tracing::info!("Docker not available: {e}"); - } - } - } - - /// Wait for the next Docker housekeeping event: a stats-poll tick while - /// the daemon is connected, or a retry tick while it is absent. Pending - /// when connected with stats polling off. - pub(super) async fn docker_tick(&mut self) -> DockerTick { - if self.docker_manager.is_some() { - match self.docker_stats_interval.as_mut() { - Some(iv) => { - iv.tick().await; - DockerTick::PollStats - } - None => std::future::pending().await, - } - } else { - self.docker_retry_interval.tick().await; - DockerTick::Retry - } - } - - /// Poll container stats once; a failing daemon demotes the runtime. - pub(super) async fn poll_docker_stats(&mut self, write: &mut S) -> anyhow::Result<()> - where - S: SinkExt + Unpin, - { - if let Some(dm) = self.docker_manager.as_mut() - && let Err(e) = dm.poll_stats().await - { - tracing::warn!("Docker stats polling failed: {e}"); - self.demote_docker(write).await?; - } - Ok(()) - } - - /// Retry the daemon connection; success re-advertises the feature. - pub(super) async fn retry_docker(&mut self, write: &mut S) -> anyhow::Result<()> - where - S: SinkExt + Unpin, - { - tracing::debug!("Retrying Docker connection..."); - match DockerManager::try_new(self.docker_tx.clone(), Arc::clone(&self.capabilities)) { - Ok(dm) => match dm.verify_connection().await { - Ok(()) => { - tracing::info!("Docker daemon now available"); - self.docker_manager = Some(dm); - self.docker_available = true; - let msg = AgentMessage::FeaturesUpdate { - features: vec!["docker".to_string()], - }; - send_msg(write, &msg).await?; - } - Err(e) => { - tracing::debug!("Docker still not available: {e}"); - } - }, - Err(e) => { - tracing::debug!("Docker still not available: {e}"); - } - } - Ok(()) - } - - async fn demote_docker(&mut self, write: &mut S) -> anyhow::Result<()> - where - S: SinkExt + Unpin, - { - if let Some(dm) = self.docker_manager.as_mut() { - dm.cleanup(); - } - self.docker_manager = None; - self.docker_stats_interval = None; - - if self.docker_available { - self.docker_available = false; - let msg = AgentMessage::FeaturesUpdate { features: vec![] }; - send_msg(write, &msg).await?; - } - - Ok(()) - } - /// Tear down every in-flight resource this connection owns. Called on all /// connection-exit paths (server close, WS error, stream end). pub(super) fn shutdown(&mut self) { @@ -254,9 +125,7 @@ impl ConnectionRuntime { self.network_prober.stop_all(); self.unlock_checker.stop(); self.file_manager.cancel_all_transfers(); - if let Some(dm) = self.docker_manager.as_mut() { - dm.cleanup(); - } + self.docker.cleanup(); } /// Parse and execute one server text frame against this connection's @@ -472,25 +341,11 @@ impl ConnectionRuntime { ) .await?; } - // --- Docker messages --- - ServerMessage::DockerStartStats { interval_secs } => { - if self.docker_manager.is_some() { - let secs = interval_secs.max(1); - tracing::info!("Starting Docker stats polling every {secs}s"); - let mut iv = tokio::time::interval(Duration::from_secs(secs as u64)); - iv.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - self.docker_stats_interval = Some(iv); - } else { - tracing::warn!("DockerStartStats received but Docker is not available"); - let unavailable = AgentMessage::DockerUnavailable { msg_id: None }; - send_msg(write, &unavailable).await?; - } - } - ServerMessage::DockerStopStats => { - tracing::info!("Stopping Docker stats polling"); - self.docker_stats_interval = None; - } - ServerMessage::DockerListContainers { .. } + // --- Docker messages --- (availability FSM + replies live in + // docker_subsystem; the reply is whatever the transition emits) + msg @ (ServerMessage::DockerStartStats { .. } + | ServerMessage::DockerStopStats + | ServerMessage::DockerListContainers { .. } | ServerMessage::DockerLogsStart { .. } | ServerMessage::DockerLogsStop { .. } | ServerMessage::DockerEventsStart @@ -498,18 +353,9 @@ impl ConnectionRuntime { | ServerMessage::DockerContainerAction { .. } | ServerMessage::DockerGetInfo { .. } | ServerMessage::DockerListNetworks { .. } - | ServerMessage::DockerListVolumes { .. } => { - if let Some(dm) = self.docker_manager.as_mut() { - if let Err(e) = dm.handle_server_message(msg.clone()).await { - tracing::warn!("Docker runtime became unavailable: {e}"); - self.demote_docker(write).await?; - } - } else { - tracing::warn!("Docker message received but Docker is not available"); - let unavailable = AgentMessage::DockerUnavailable { - msg_id: docker_request_msg_id(&msg), - }; - send_msg(write, &unavailable).await?; + | ServerMessage::DockerListVolumes { .. }) => { + if let Some(reply) = self.docker.handle(msg).await { + send_msg(write, &reply).await?; } } // Firewall blocklist variants — dispatched to the FirewallManager @@ -594,17 +440,6 @@ fn spawn_capability_denied( }); } -fn docker_request_msg_id(msg: &ServerMessage) -> Option { - match msg { - ServerMessage::DockerListContainers { msg_id } - | ServerMessage::DockerGetInfo { msg_id } - | ServerMessage::DockerListNetworks { msg_id } - | ServerMessage::DockerListVolumes { msg_id } => Some(msg_id.clone()), - ServerMessage::DockerContainerAction { msg_id, .. } => Some(msg_id.clone()), - _ => None, - } -} - async fn execute_command( task_id: &str, command: &str, @@ -662,78 +497,6 @@ mod tests { use crate::firewall::nft::CliNftExecutor; use serverbee_common::constants::CapabilityDeniedReason; - #[test] - fn test_docker_request_msg_id_extracts_only_request_variants() { - // Variants that carry a msg_id return it... - assert_eq!( - docker_request_msg_id(&ServerMessage::DockerListContainers { - msg_id: "a".to_string() - }), - Some("a".to_string()) - ); - assert_eq!( - docker_request_msg_id(&ServerMessage::DockerGetInfo { - msg_id: "b".to_string() - }), - Some("b".to_string()) - ); - assert_eq!( - docker_request_msg_id(&ServerMessage::DockerListNetworks { - msg_id: "c".to_string() - }), - Some("c".to_string()) - ); - assert_eq!( - docker_request_msg_id(&ServerMessage::DockerListVolumes { - msg_id: "d".to_string() - }), - Some("d".to_string()) - ); - assert_eq!( - docker_request_msg_id(&ServerMessage::DockerContainerAction { - msg_id: "e".to_string(), - container_id: "cid".to_string(), - action: serverbee_common::docker_types::DockerAction::Restart { timeout: None }, - }), - Some("e".to_string()) - ); - // ...non-request docker variants return None. - assert_eq!(docker_request_msg_id(&ServerMessage::DockerStopStats), None); - assert_eq!(docker_request_msg_id(&ServerMessage::Ping), None); - } - - #[test] - fn test_docker_request_msg_id_none_for_log_and_event_variants() { - // Streaming control variants carry no request msg_id -> None. - assert_eq!( - docker_request_msg_id(&ServerMessage::DockerLogsStart { - session_id: "s".to_string(), - container_id: "c".to_string(), - tail: None, - follow: false, - }), - None - ); - assert_eq!( - docker_request_msg_id(&ServerMessage::DockerLogsStop { - session_id: "s".to_string(), - }), - None - ); - assert_eq!( - docker_request_msg_id(&ServerMessage::DockerEventsStart), - None - ); - assert_eq!( - docker_request_msg_id(&ServerMessage::DockerEventsStop), - None - ); - assert_eq!( - docker_request_msg_id(&ServerMessage::DockerStartStats { interval_secs: 5 }), - None - ); - } - // ---------------------------------------------------------------------- // `execute_command` — pure-ish process helper. These shell out to `sh`, // which is always present on macOS/Linux CI, so they remain deterministic. @@ -1476,20 +1239,20 @@ mod tests { &msgs[0], AgentMessage::DockerUnavailable { msg_id: None } )); - assert!(h.runtime.docker_stats_interval.is_none()); + assert!(!h.runtime.docker.stats_polling_active()); } #[tokio::test] async fn test_dispatch_docker_stop_stats_clears_interval() { let mut h = Harness::new(ALL_CAPS, FileConfig::default()); // Pre-seed an interval so we can observe it being cleared. - h.runtime.docker_stats_interval = Some(tokio::time::interval(Duration::from_secs(60))); + h.runtime.docker.arm_stats_polling_for_test(); let mut sink = RecordingSink::new(); h.dispatch(r#"{"type":"docker_stop_stats"}"#, &mut sink) .await .unwrap(); assert_eq!(sink.sent_count(), 0); - assert!(h.runtime.docker_stats_interval.is_none()); + assert!(!h.runtime.docker.stats_polling_active()); } #[tokio::test] From ce817911c535e4d5047f770b6b0d88c2ba52f21a Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sun, 12 Jul 2026 00:28:43 +0800 Subject: [PATCH 35/44] refactor(agent): put a MetricsSource seam behind report assembly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Collector hardwired sysinfo handles and eleven free-function reads, so the assembly logic (network rate differencing, counter-reset saturation, the elapsed clamp, temperature/GPU gating) could only run against the real host and tests degraded to loose >=0 invariants. The MetricsSource trait is the seam: SysinfoSource is the production adapter wrapping the existing submodules (which are unchanged), and the test FakeSource drives every assembly path with arbitrary readings — counter resets now provably read as zero-speed samples, sub-second windows provably clamp, disabled sensors provably gate to None. Real-host smoke coverage in collector/tests.rs is preserved. --- crates/agent/src/collector/mod.rs | 331 ++++++++++++++++++--------- crates/agent/src/collector/source.rs | 145 ++++++++++++ crates/agent/src/collector/tests.rs | 31 +++ 3 files changed, 403 insertions(+), 104 deletions(-) create mode 100644 crates/agent/src/collector/source.rs diff --git a/crates/agent/src/collector/mod.rs b/crates/agent/src/collector/mod.rs index 2d46f5556..123067aae 100644 --- a/crates/agent/src/collector/mod.rs +++ b/crates/agent/src/collector/mod.rs @@ -6,18 +6,23 @@ mod load; mod memory; mod network; mod process; +mod source; mod temperature; pub mod virtualization; use std::time::Instant; use serverbee_common::types::{SystemInfo, SystemReport}; -use sysinfo::{Networks, ProcessRefreshKind, ProcessesToUpdate, System}; -pub struct Collector { - sys: System, - networks: Networks, - prev_disk_io: std::collections::HashMap, +use source::{MetricsSource, SysinfoSource}; + +/// Report assembly over a [`MetricsSource`]: rate differencing for the +/// cumulative network counters, the elapsed-window guard, and the +/// temperature/GPU enable gating live here — the source only reads the host. +/// Production code uses the default `SysinfoSource`; tests inject a fake to +/// drive every assembly path with arbitrary readings. +pub struct Collector { + source: S, prev_net_in: u64, prev_net_out: u64, prev_time: Instant, @@ -25,16 +30,21 @@ pub struct Collector { enable_gpu: bool, } -impl Collector { +impl Collector { pub fn new(enable_temperature: bool, enable_gpu: bool) -> Self { - let mut sys = System::new_all(); - sys.refresh_all(); - let networks = Networks::new_with_refreshed_list(); - let (net_in, net_out) = network::total_bytes(&networks); + Self::with_source(SysinfoSource::new(), enable_temperature, enable_gpu) + } + + pub fn system_info(&self) -> SystemInfo { + self.source.system_info() + } +} + +impl Collector { + fn with_source(source: S, enable_temperature: bool, enable_gpu: bool) -> Self { + let (net_in, net_out) = source.net_total_bytes(); Self { - sys, - networks, - prev_disk_io: std::collections::HashMap::new(), + source, prev_net_in: net_in, prev_net_out: net_out, prev_time: Instant::now(), @@ -44,140 +54,253 @@ impl Collector { } pub fn collect(&mut self) -> SystemReport { - self.sys.refresh_cpu_usage(); - self.sys.refresh_memory(); - self.sys.refresh_processes_specifics( - ProcessesToUpdate::All, - true, - ProcessRefreshKind::nothing(), - ); - self.networks.refresh(true); - - let elapsed = self.prev_time.elapsed().as_secs_f64().max(1.0); - let (net_in, net_out) = network::total_bytes(&self.networks); + self.source.refresh(); + let elapsed = self.prev_time.elapsed().as_secs_f64(); + self.prev_time = Instant::now(); + self.assemble_report(elapsed) + } + + /// Fold one refreshed sample into a report. Split from [`Self::collect`] + /// so tests control the elapsed window directly. + fn assemble_report(&mut self, elapsed: f64) -> SystemReport { + // Guard the rate denominator: a sub-second (or clock-skewed) window + // must not inflate speeds or divide by zero. + let elapsed = elapsed.max(1.0); + + // Cumulative counters difference into per-second rates; saturating_sub + // absorbs counter resets (interface re-enumeration, rollover) as a + // zero-speed sample instead of a negative one. + let (net_in, net_out) = self.source.net_total_bytes(); let net_in_speed = (net_in.saturating_sub(self.prev_net_in) as f64 / elapsed) as i64; let net_out_speed = (net_out.saturating_sub(self.prev_net_out) as f64 / elapsed) as i64; self.prev_net_in = net_in; self.prev_net_out = net_out; - self.prev_time = Instant::now(); - let disk_io = disk_io::collect(elapsed, &mut self.prev_disk_io); + let disk_io = self.source.disk_io(elapsed); let temperature = if self.enable_temperature { - temperature::get_temperature() + self.source.temperature() } else { None }; + let (load1, load5, load15) = self.source.load_averages(); + SystemReport { - cpu: cpu::usage(&self.sys), - mem_used: memory::mem_used(&self.sys), - swap_used: memory::swap_used(&self.sys), - disk_used: disk::used(), + cpu: self.source.cpu_usage(), + mem_used: self.source.mem_used(), + swap_used: self.source.swap_used(), + disk_used: self.source.disk_used(), net_in_speed, net_out_speed, net_in_transfer: net_in as i64, net_out_transfer: net_out as i64, - load1: load::load1(), - load5: load::load5(), - load15: load::load15(), - tcp_conn: process::tcp_connections(), - udp_conn: process::udp_connections(), - process_count: process::count(&self.sys), - uptime: System::uptime(), + load1, + load5, + load15, + tcp_conn: self.source.tcp_connections(), + udp_conn: self.source.udp_connections(), + process_count: self.source.process_count(), + uptime: self.source.uptime(), disk_io, temperature, gpu: if self.enable_gpu { - gpu::get_gpu_report() + self.source.gpu() } else { None }, } } - - pub fn system_info(&self) -> SystemInfo { - SystemInfo { - protocol_version: 0, - cpu_name: cpu::name(&self.sys), - cpu_cores: cpu::cores(&self.sys), - cpu_arch: cpu::arch(), - os: System::long_os_version().unwrap_or_default(), - kernel_version: System::kernel_version().unwrap_or_default(), - mem_total: memory::mem_total(&self.sys), - swap_total: memory::swap_total(&self.sys), - disk_total: disk::total(), - ipv4: None, - ipv6: None, - virtualization: virtualization::detect(), - agent_version: serverbee_common::constants::VERSION.to_string(), - features: Vec::new(), - } - } } #[cfg(test)] mod tests; #[cfg(test)] -mod branch_tests { +mod assembly_tests { + use super::source::MetricsSource; use super::*; + use serverbee_common::types::{DiskIo, GpuInfo, GpuReport}; + + /// Test adapter: every reading is a plain field, so assembly paths can be + /// driven with arbitrary host states. + struct FakeSource { + cpu: f64, + mem_used: i64, + swap_used: i64, + disk_used: i64, + net_totals: (u64, u64), + loads: (f64, f64, f64), + tcp: i32, + udp: i32, + processes: i32, + uptime: u64, + disk_io: Option>, + temperature: Option, + gpu: Option, + refreshes: u32, + } + + impl Default for FakeSource { + fn default() -> Self { + Self { + cpu: 12.5, + mem_used: 1024, + swap_used: 256, + disk_used: 4096, + net_totals: (10_000, 20_000), + loads: (0.5, 0.4, 0.3), + tcp: 7, + udp: 3, + processes: 42, + uptime: 3600, + disk_io: None, + temperature: Some(55.0), + gpu: Some(GpuReport { + count: 1, + average_usage: 30.0, + detailed_info: vec![GpuInfo { + name: "FakeGPU".to_string(), + mem_total: 8000, + mem_used: 2000, + utilization: 30.0, + temperature: 60.0, + }], + }), + refreshes: 0, + } + } + } + + impl MetricsSource for FakeSource { + fn refresh(&mut self) { + self.refreshes += 1; + } + fn cpu_usage(&self) -> f64 { + self.cpu + } + fn mem_used(&self) -> i64 { + self.mem_used + } + fn swap_used(&self) -> i64 { + self.swap_used + } + fn disk_used(&self) -> i64 { + self.disk_used + } + fn net_total_bytes(&self) -> (u64, u64) { + self.net_totals + } + fn load_averages(&self) -> (f64, f64, f64) { + self.loads + } + fn tcp_connections(&self) -> i32 { + self.tcp + } + fn udp_connections(&self) -> i32 { + self.udp + } + fn process_count(&self) -> i32 { + self.processes + } + fn uptime(&self) -> u64 { + self.uptime + } + fn disk_io(&mut self, _elapsed: f64) -> Option> { + self.disk_io.clone() + } + fn temperature(&self) -> Option { + self.temperature + } + fn gpu(&self) -> Option { + self.gpu.clone() + } + } + + fn make_collector(source: FakeSource) -> Collector { + Collector::with_source(source, true, true) + } + + #[test] + fn net_speeds_difference_cumulative_counters_over_elapsed() { + // prev seeded from the constructor reading (10k/20k); the next sample + // adds 30k/60k over a 10s window -> 3k/6k per second. + let mut c = make_collector(FakeSource::default()); + c.source.net_totals = (40_000, 80_000); + let report = c.assemble_report(10.0); + assert_eq!(report.net_in_speed, 3_000); + assert_eq!(report.net_out_speed, 6_000); + assert_eq!(report.net_in_transfer, 40_000); + assert_eq!(report.net_out_transfer, 80_000); + } + + #[test] + fn counter_reset_saturates_to_zero_speed() { + // A counter that went backwards (interface re-enumeration) must read + // as a zero-speed sample, never a negative rate. + let mut c = make_collector(FakeSource::default()); + c.source.net_totals = (500, 700); + let report = c.assemble_report(5.0); + assert_eq!(report.net_in_speed, 0); + assert_eq!(report.net_out_speed, 0); + // The new (lower) totals still become the baseline for the next window. + c.source.net_totals = (1_500, 1_700); + let next = c.assemble_report(1.0); + assert_eq!(next.net_in_speed, 1_000); + assert_eq!(next.net_out_speed, 1_000); + } + + #[test] + fn sub_second_elapsed_is_clamped_to_one_second() { + // elapsed 0.0 (or any sub-second window) divides by 1.0: finite, + // non-inflated speeds. + let mut c = make_collector(FakeSource::default()); + c.source.net_totals = (10_100, 20_200); + let report = c.assemble_report(0.0); + assert_eq!(report.net_in_speed, 100); + assert_eq!(report.net_out_speed, 200); + } #[test] - fn collect_with_temperature_disabled_yields_none() { - // When `enable_temperature` is false the collector skips the sensor read - // entirely and reports `temperature: None` (the disabled branch in - // `Collector::collect`). - let mut collector = Collector::new(false, false); - let report = collector.collect(); + fn disabled_temperature_and_gpu_are_gated_to_none() { + // The source has readings, but disabled collection must report None — + // the gate lives in assembly, not in the source. + let mut c = Collector::with_source(FakeSource::default(), false, false); + let report = c.assemble_report(1.0); assert!(report.temperature.is_none()); + assert!(report.gpu.is_none()); } #[test] - fn collect_with_gpu_enabled_does_not_panic() { - // With `enable_gpu` true the collector invokes `gpu::get_gpu_report()`. - // Without the `gpu` cargo feature (the default, and the case on CI hosts - // with no NVIDIA GPU) this returns None, but the branch must be reached - // and must not panic. When a report is present, validate its invariants. - let mut collector = Collector::new(false, true); - let report = collector.collect(); - if let Some(gpu) = report.gpu { - assert_eq!(gpu.count as usize, gpu.detailed_info.len()); - assert!(gpu.average_usage.is_finite()); - } + fn enabled_temperature_and_gpu_pass_through() { + let mut c = make_collector(FakeSource::default()); + let report = c.assemble_report(1.0); + assert_eq!(report.temperature, Some(55.0)); + assert_eq!(report.gpu.as_ref().map(|g| g.count), Some(1)); } #[test] - fn collect_report_field_invariants() { - // Cross-field invariants on a freshly collected report: connection and - // network counters are non-negative and load averages are finite. - let mut collector = Collector::new(false, false); - let report = collector.collect(); - assert!(report.tcp_conn >= 0); - assert!(report.udp_conn >= 0); - assert!(report.net_in_transfer >= 0); - assert!(report.net_out_transfer >= 0); - assert!(report.net_in_speed >= 0); - assert!(report.net_out_speed >= 0); - assert!(report.load1 >= 0.0 && report.load1.is_finite()); - assert!(report.load5 >= 0.0 && report.load5.is_finite()); - assert!(report.load15 >= 0.0 && report.load15.is_finite()); - assert!(report.swap_used >= 0); + fn gauges_pass_through_untransformed() { + let mut c = make_collector(FakeSource::default()); + let report = c.assemble_report(1.0); + assert!((report.cpu - 12.5).abs() < f64::EPSILON); + assert_eq!(report.mem_used, 1024); + assert_eq!(report.swap_used, 256); + assert_eq!(report.disk_used, 4096); + assert!((report.load1 - 0.5).abs() < f64::EPSILON); + assert!((report.load15 - 0.3).abs() < f64::EPSILON); + assert_eq!(report.tcp_conn, 7); + assert_eq!(report.udp_conn, 3); + assert_eq!(report.process_count, 42); + assert_eq!(report.uptime, 3600); } #[test] - fn system_info_static_assembly_fields() { - // Assert the `system_info` assembly fields not covered by the existing - // populated-fields test (arch, agent version, protocol version, and the - // default-empty optional/ip/feature fields). - let collector = Collector::new(false, false); - let info = collector.system_info(); - assert!(!info.cpu_arch.is_empty()); - assert!(!info.agent_version.is_empty()); - assert_eq!(info.protocol_version, 0); - assert!(info.ipv4.is_none()); - assert!(info.ipv6.is_none()); - assert!(info.features.is_empty()); - assert!(info.swap_total >= 0); + fn collect_refreshes_the_source_before_sampling() { + let mut c = make_collector(FakeSource::default()); + let _ = c.collect(); + let _ = c.collect(); + assert_eq!(c.source.refreshes, 2); } } diff --git a/crates/agent/src/collector/source.rs b/crates/agent/src/collector/source.rs new file mode 100644 index 000000000..f745a5aa5 --- /dev/null +++ b/crates/agent/src/collector/source.rs @@ -0,0 +1,145 @@ +//! The host-reading seam behind report assembly. +//! +//! [`MetricsSource`] is everything `Collector` needs to read from the host; +//! [`SysinfoSource`] is the production adapter wrapping the sysinfo handles +//! and the per-metric submodules. Tests drive the assembly logic (rate +//! differencing, counter-reset saturation, the elapsed guard, enable gating) +//! through a fake source instead of the real machine. + +use std::collections::HashMap; + +use serverbee_common::types::{DiskIo, GpuReport, SystemInfo}; +use sysinfo::{Networks, ProcessRefreshKind, ProcessesToUpdate, System}; + +use super::{cpu, disk, disk_io, load, memory, network, process, temperature, virtualization}; + +/// One connection's window onto the host: raw gauges, cumulative counters, +/// and windowed per-device I/O rates. Implementations own whatever handles +/// and previous-sample state the readings require. +pub(crate) trait MetricsSource { + /// Refresh whatever underlying handles need refreshing before a sample. + fn refresh(&mut self); + fn cpu_usage(&self) -> f64; + fn mem_used(&self) -> i64; + fn swap_used(&self) -> i64; + fn disk_used(&self) -> i64; + /// Cumulative interface byte counters since boot (in, out). + fn net_total_bytes(&self) -> (u64, u64); + /// (load1, load5, load15) + fn load_averages(&self) -> (f64, f64, f64); + fn tcp_connections(&self) -> i32; + fn udp_connections(&self) -> i32; + fn process_count(&self) -> i32; + fn uptime(&self) -> u64; + /// Per-device I/O rates over the elapsed window, when the platform + /// exposes them. + fn disk_io(&mut self, elapsed: f64) -> Option>; + fn temperature(&self) -> Option; + fn gpu(&self) -> Option; +} + +/// Production adapter: sysinfo handles plus the per-metric submodules. +pub(crate) struct SysinfoSource { + sys: System, + networks: Networks, + prev_disk_io: HashMap, +} + +impl SysinfoSource { + pub(crate) fn new() -> Self { + let mut sys = System::new_all(); + sys.refresh_all(); + let networks = Networks::new_with_refreshed_list(); + Self { + sys, + networks, + prev_disk_io: HashMap::new(), + } + } + + /// Static host facts for `SystemInfo`; not part of the [`MetricsSource`] + /// seam because nothing in the assembly transforms them. + pub(crate) fn system_info(&self) -> SystemInfo { + SystemInfo { + protocol_version: 0, + cpu_name: cpu::name(&self.sys), + cpu_cores: cpu::cores(&self.sys), + cpu_arch: cpu::arch(), + os: System::long_os_version().unwrap_or_default(), + kernel_version: System::kernel_version().unwrap_or_default(), + mem_total: memory::mem_total(&self.sys), + swap_total: memory::swap_total(&self.sys), + disk_total: disk::total(), + ipv4: None, + ipv6: None, + virtualization: virtualization::detect(), + agent_version: serverbee_common::constants::VERSION.to_string(), + features: Vec::new(), + } + } +} + +impl MetricsSource for SysinfoSource { + fn refresh(&mut self) { + self.sys.refresh_cpu_usage(); + self.sys.refresh_memory(); + self.sys.refresh_processes_specifics( + ProcessesToUpdate::All, + true, + ProcessRefreshKind::nothing(), + ); + self.networks.refresh(true); + } + + fn cpu_usage(&self) -> f64 { + cpu::usage(&self.sys) + } + + fn mem_used(&self) -> i64 { + memory::mem_used(&self.sys) + } + + fn swap_used(&self) -> i64 { + memory::swap_used(&self.sys) + } + + fn disk_used(&self) -> i64 { + disk::used() + } + + fn net_total_bytes(&self) -> (u64, u64) { + network::total_bytes(&self.networks) + } + + fn load_averages(&self) -> (f64, f64, f64) { + (load::load1(), load::load5(), load::load15()) + } + + fn tcp_connections(&self) -> i32 { + process::tcp_connections() + } + + fn udp_connections(&self) -> i32 { + process::udp_connections() + } + + fn process_count(&self) -> i32 { + process::count(&self.sys) + } + + fn uptime(&self) -> u64 { + System::uptime() + } + + fn disk_io(&mut self, elapsed: f64) -> Option> { + disk_io::collect(elapsed, &mut self.prev_disk_io) + } + + fn temperature(&self) -> Option { + temperature::get_temperature() + } + + fn gpu(&self) -> Option { + super::gpu::get_gpu_report() + } +} diff --git a/crates/agent/src/collector/tests.rs b/crates/agent/src/collector/tests.rs index 54bda23a8..1080c3813 100644 --- a/crates/agent/src/collector/tests.rs +++ b/crates/agent/src/collector/tests.rs @@ -62,3 +62,34 @@ fn test_collect_disk_io_first_sample_is_empty_on_non_linux() { let report = collector.collect(); assert_eq!(report.disk_io, Some(vec![])); } + +#[test] +fn test_collect_with_gpu_enabled_does_not_panic() { + // With `enable_gpu` true the collector invokes the real GPU probe through + // SysinfoSource. Without the `gpu` cargo feature (the default, and the + // case on CI hosts with no NVIDIA GPU) this returns None, but the branch + // must be reached and must not panic. When a report is present, validate + // its invariants. + let mut collector = Collector::new(false, true); + let report = collector.collect(); + if let Some(gpu) = report.gpu { + assert_eq!(gpu.count as usize, gpu.detailed_info.len()); + assert!(gpu.average_usage.is_finite()); + } +} + +#[test] +fn test_system_info_static_assembly_fields() { + // The `system_info` assembly fields not covered by the populated-fields + // test: arch, agent version, protocol version, and the default-empty + // optional/ip/feature fields. + let collector = Collector::new(false, false); + let info = collector.system_info(); + assert!(!info.cpu_arch.is_empty()); + assert!(!info.agent_version.is_empty()); + assert_eq!(info.protocol_version, 0); + assert!(info.ipv4.is_none()); + assert!(info.ipv6.is_none()); + assert!(info.features.is_empty()); + assert!(info.swap_total >= 0); +} From 864eed8a919a10fe253340c17f841a59a8709689 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sun, 12 Jul 2026 00:34:46 +0800 Subject: [PATCH 36/44] refactor(server): funnel effective-capability checks through capability_gate The security-sensitive priority rule (live agent report first, persisted mirror before the first SystemInfo) was encoded three times: in agent_manager::capability_denied_reason, in security.rs's private effective_caps_or_db, and as a cache-only ad-hoc check in the ip-quality run-now guard that hard-denied during the pre-report window. Name the rule once (agent_manager::effective_capabilities_or), expose the raw mask through capability_gate::effective_capabilities for consumers with bespoke error shapes, and route both bypasses through it. The ip-quality guard now follows the same mirror-fallback semantics as every other gate instead of its own stricter variant. --- crates/server/src/router/api/ip_quality.rs | 7 ++-- crates/server/src/router/ws/agent/security.rs | 27 +++------------ crates/server/src/service/agent_manager.rs | 21 +++++++----- crates/server/src/service/capability_gate.rs | 33 +++++++++++++++++++ 4 files changed, 52 insertions(+), 36 deletions(-) diff --git a/crates/server/src/router/api/ip_quality.rs b/crates/server/src/router/api/ip_quality.rs index b2ef19cdc..91bacd9b2 100644 --- a/crates/server/src/router/api/ip_quality.rs +++ b/crates/server/src/router/api/ip_quality.rs @@ -273,11 +273,8 @@ async fn check_server( // capability. The agent would silently ignore the message, giving the UI // false success. Capabilities are agent-owned, so the only fix is to edit // the agent's config file (or CLI flags) on the host. - let agent_has = state - .agent_manager - .get_agent_local_capabilities(&id) - .map(|caps| has_capability(caps, CAP_IP_QUALITY)) - .unwrap_or(false); + let caps = crate::service::capability_gate::effective_capabilities(&state, &id).await; + let agent_has = has_capability(caps, CAP_IP_QUALITY); if !agent_has { return Err(AppError::Conflict( diff --git a/crates/server/src/router/ws/agent/security.rs b/crates/server/src/router/ws/agent/security.rs index dcd29c594..446763d85 100644 --- a/crates/server/src/router/ws/agent/security.rs +++ b/crates/server/src/router/ws/agent/security.rs @@ -24,35 +24,16 @@ fn is_high_risk_cap(cap: &str) -> bool { matches!(cap, "terminal" | "exec" | "file" | "docker") } -/// Effective capabilities for inbound-data gating. The hot-path cache is -/// populated once the agent sends `SystemInfo` on the current connection; -/// before that, fall back to the DB row. -async fn effective_caps_or_db(state: &Arc, server_id: &str) -> u32 { - match state.agent_manager.get_effective_capabilities(server_id) { - Some(c) => c, - None => { - use crate::entity::server; - use sea_orm::EntityTrait; - server::Entity::find_by_id(server_id) - .one(&state.db) - .await - .ok() - .flatten() - .and_then(|s| u32::try_from(s.capabilities).ok()) - .unwrap_or(0) - } - } -} - -/// Re-check `cap` before accepting inbound agent data. On denial, writes a -/// `denied_action` audit row and returns false. +/// Re-check `cap` before accepting inbound agent data (the priority rule — +/// live report first, mirror fallback — lives in capability_gate). On denial, +/// writes a `denied_action` audit row and returns false. async fn gate_inbound_data( state: &Arc, server_id: &str, cap: u32, denied_action: &str, ) -> bool { - let caps = effective_caps_or_db(state, server_id).await; + let caps = crate::service::capability_gate::effective_capabilities(state, server_id).await; if has_capability(caps, cap) { return true; } diff --git a/crates/server/src/service/agent_manager.rs b/crates/server/src/service/agent_manager.rs index 68d6087c7..26597440c 100644 --- a/crates/server/src/service/agent_manager.rs +++ b/crates/server/src/service/agent_manager.rs @@ -631,21 +631,26 @@ impl AgentManager { self.get_agent_local_capabilities(server_id) } + /// The one encoding of the capability priority rule: the live + /// agent-reported bitmask decides, and `mirror_caps` (the persisted + /// `servers.capabilities` mirror) fills the brief window before the + /// agent's first `SystemInfo` (or while it is offline). Every "is this + /// capability effective" answer must resolve through here. + pub fn effective_capabilities_or(&self, server_id: &str, mirror_caps: u32) -> u32 { + self.get_agent_local_capabilities(server_id) + .unwrap_or(mirror_caps) + } + /// Returns `Some(reason)` when `cap_bit` is NOT available on the agent, or - /// `None` when it is allowed. `mirror_caps` is the persisted - /// `servers.capabilities` mirror, used as a fallback for the brief window - /// before the agent's first `SystemInfo` (or while it is offline). The - /// reason is always agent-side: the server has no say in capabilities. + /// `None` when it is allowed. The reason is always agent-side: the server + /// has no say in capabilities. pub fn capability_denied_reason( &self, server_id: &str, mirror_caps: u32, cap_bit: u32, ) -> Option<&'static str> { - let caps = self - .get_agent_local_capabilities(server_id) - .unwrap_or(mirror_caps); - if has_capability(caps, cap_bit) { + if has_capability(self.effective_capabilities_or(server_id, mirror_caps), cap_bit) { None } else { Some("agent_capability_disabled") diff --git a/crates/server/src/service/capability_gate.rs b/crates/server/src/service/capability_gate.rs index 14ac2c076..d22001dd9 100644 --- a/crates/server/src/service/capability_gate.rs +++ b/crates/server/src/service/capability_gate.rs @@ -13,6 +13,25 @@ use crate::service::audit::AuditService; use crate::service::server::ServerService; use crate::state::AppState; +/// The raw agent-effective bitmask for consumers that need the mask itself +/// (inbound-data gates, availability checks with bespoke error shapes) rather +/// than a deny decision: live agent report first, persisted mirror before the +/// first `SystemInfo`, zero when the server row is gone. Same priority rule +/// as [`require_capability`], resolved through the same funnel. +pub async fn effective_capabilities(state: &AppState, server_id: &str) -> u32 { + // Hot path: a reported bitmask decides without touching the DB. + if let Some(caps) = state.agent_manager.get_effective_capabilities(server_id) { + return caps; + } + let mirror = ServerService::get_server(&state.db, server_id) + .await + .map(|s| s.capabilities as u32) + .unwrap_or(0); + state + .agent_manager + .effective_capabilities_or(server_id, mirror) +} + /// Resolve the server row and apply the agent-owned capability policy. /// Returns the server model on success so callers don't re-query. pub async fn require_capability( @@ -106,6 +125,20 @@ mod tests { SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080) } + #[tokio::test] + async fn effective_capabilities_prefers_report_then_mirror_then_zero() { + let (state, _tmp) = setup_state_with_server(CAP_FILE).await; + // No report yet: the persisted mirror decides. + assert_eq!(effective_capabilities(&state, "srv-1").await, CAP_FILE); + // A live report overrides the mirror entirely. + state + .agent_manager + .update_agent_local_capabilities("srv-1", CAP_DOCKER); + assert_eq!(effective_capabilities(&state, "srv-1").await, CAP_DOCKER); + // Unknown server resolves to no capabilities, not an error. + assert_eq!(effective_capabilities(&state, "nope").await, 0); + } + #[tokio::test] async fn missing_server_is_not_found() { let (state, _tmp) = setup_state_with_server(CAP_FILE).await; From 5c1d03d7a97b9685c6692a6d4a9a157e159af1f3 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sun, 12 Jul 2026 00:39:17 +0800 Subject: [PATCH 37/44] refactor(web): generate capability bit definitions from ALL_CAPABILITIES The web capability constants were a hand-typed mirror of the Rust table, synced by comments, with CAP_DEFAULT = 1852 duplicated as a magic number. capability-bits.generated.ts is now emitted from ALL_CAPABILITIES by a cargo example (same pipeline shape as dump_openapi -> api-types), the generator asserts CAP_DEFAULT equals the OR of default_enabled bits, and capabilities.ts re-exports the generated definitions alongside its hand-written helpers so import sites are unchanged. Adding a capability bit is now one Rust table row plus regeneration (iOS remains a deliberate cross-repo hand mirror). --- apps/web/package.json | 3 +- apps/web/src/lib/capabilities.ts | 71 ++++++------------- apps/web/src/lib/capability-bits.generated.ts | 33 +++++++++ .../common/examples/dump_capabilities_ts.rs | 57 +++++++++++++++ 4 files changed, 113 insertions(+), 51 deletions(-) create mode 100644 apps/web/src/lib/capability-bits.generated.ts create mode 100644 crates/common/examples/dump_capabilities_ts.rs diff --git a/apps/web/package.json b/apps/web/package.json index 6f46b96d9..823be49ad 100644 --- a/apps/web/package.json +++ b/apps/web/package.json @@ -11,7 +11,8 @@ "preview": "vite preview", "test": "vitest run", "test:watch": "vitest", - "generate:api-types": "cargo run --manifest-path ../../Cargo.toml --example dump_openapi > openapi.json && npx openapi-typescript openapi.json -o src/lib/api-types.ts" + "generate:api-types": "cargo run --manifest-path ../../Cargo.toml --example dump_openapi > openapi.json && npx openapi-typescript openapi.json -o src/lib/api-types.ts", + "generate:capabilities": "cargo run --manifest-path ../../Cargo.toml --example dump_capabilities_ts > src/lib/capability-bits.generated.ts" }, "dependencies": { "@base-ui/react": "^1.2.0", diff --git a/apps/web/src/lib/capabilities.ts b/apps/web/src/lib/capabilities.ts index d763d20fa..1d2005fc9 100644 --- a/apps/web/src/lib/capabilities.ts +++ b/apps/web/src/lib/capabilities.ts @@ -1,54 +1,25 @@ -export const CAP_TERMINAL = 1 -export const CAP_EXEC = 2 -export const CAP_UPGRADE = 4 -export const CAP_PING_ICMP = 8 -export const CAP_PING_TCP = 16 -export const CAP_PING_HTTP = 32 -export const CAP_FILE = 64 -export const CAP_DOCKER = 128 -export const CAP_SECURITY_EVENTS = 256 -export const CAP_FIREWALL_BLOCK = 512 -export const CAP_IP_QUALITY = 1024 +// Bit definitions and metadata are generated from the Rust +// ALL_CAPABILITIES table (bun run generate:capabilities); the helpers +// below are hand-written web logic on top of them. +import { CAP_DEFAULT, CAPABILITIES, type CapabilityRisk } from './capability-bits.generated' -// Mirrors CAP_DEFAULT in crates/common/src/constants.rs (1852): -// upgrade + ICMP/TCP/HTTP ping + security events + firewall blocklist + IP quality. -export const CAP_DEFAULT = 1852 - -export const CAPABILITIES = [ - { bit: CAP_TERMINAL, key: 'terminal', labelKey: 'cap_terminal' as const, risk: 'high' as const }, - { bit: CAP_EXEC, key: 'exec', labelKey: 'cap_exec' as const, risk: 'high' as const }, - { bit: CAP_UPGRADE, key: 'upgrade', labelKey: 'cap_upgrade' as const, risk: 'low' as const }, - { bit: CAP_PING_ICMP, key: 'ping_icmp', labelKey: 'cap_ping_icmp' as const, risk: 'low' as const }, - { bit: CAP_PING_TCP, key: 'ping_tcp', labelKey: 'cap_ping_tcp' as const, risk: 'low' as const }, - { bit: CAP_PING_HTTP, key: 'ping_http', labelKey: 'cap_ping_http' as const, risk: 'low' as const }, - { bit: CAP_FILE, key: 'file', labelKey: 'cap_file' as const, risk: 'high' as const }, - { bit: CAP_DOCKER, key: 'docker', labelKey: 'cap_docker' as const, risk: 'high' as const }, - { - bit: CAP_SECURITY_EVENTS, - key: 'security_events', - labelKey: 'cap_security_events' as const, - risk: 'low' as const - }, - { - bit: CAP_FIREWALL_BLOCK, - key: 'firewall_block', - labelKey: 'cap_firewall_block' as const, - // Medium, not high: the agent can only add/remove IPs in its own dedicated - // nft blocklist set (see crates/agent/src/firewall) — it can't exec code, - // read files, or flush the host firewall, and guardrails reject self-lockout - // ranges. Mirrors the Rust risk_level in crates/common/src/constants.rs. - risk: 'medium' as const - }, - { - bit: CAP_IP_QUALITY, - key: 'ip_quality', - labelKey: 'cap_ip_quality' as const, - // Mirrors the Rust risk_level in crates/common/src/constants.rs. - risk: 'medium' as const - } -] as const - -export type CapabilityRisk = (typeof CAPABILITIES)[number]['risk'] +// biome-ignore lint/performance/noBarrelFile: capabilities.ts stays the single public surface — the generated bit definitions re-export through it so import sites don't need to know about codegen +export { + CAP_DEFAULT, + CAP_DOCKER, + CAP_EXEC, + CAP_FILE, + CAP_FIREWALL_BLOCK, + CAP_IP_QUALITY, + CAP_PING_HTTP, + CAP_PING_ICMP, + CAP_PING_TCP, + CAP_SECURITY_EVENTS, + CAP_TERMINAL, + CAP_UPGRADE, + CAPABILITIES, + type CapabilityRisk +} from './capability-bits.generated' // i18n keys (servers namespace) for each risk tier's badge label. Single source // so the settings table and the capabilities dialog stay in sync. diff --git a/apps/web/src/lib/capability-bits.generated.ts b/apps/web/src/lib/capability-bits.generated.ts new file mode 100644 index 000000000..9e5d01133 --- /dev/null +++ b/apps/web/src/lib/capability-bits.generated.ts @@ -0,0 +1,33 @@ +// AUTO-GENERATED from crates/common/src/constants.rs (ALL_CAPABILITIES). +// Regenerate with `bun run generate:capabilities` — do not edit by hand. + +export const CAP_TERMINAL = 1 +export const CAP_EXEC = 2 +export const CAP_UPGRADE = 4 +export const CAP_PING_ICMP = 8 +export const CAP_PING_TCP = 16 +export const CAP_PING_HTTP = 32 +export const CAP_FILE = 64 +export const CAP_DOCKER = 128 +export const CAP_SECURITY_EVENTS = 256 +export const CAP_FIREWALL_BLOCK = 512 +export const CAP_IP_QUALITY = 1024 + +// The OR of every default_enabled bit: upgrade + ping_icmp + ping_tcp + ping_http + security_events + firewall_block + ip_quality. +export const CAP_DEFAULT = 1852 + +export const CAPABILITIES = [ + { bit: CAP_TERMINAL, key: 'terminal', labelKey: 'cap_terminal', risk: 'high' }, + { bit: CAP_EXEC, key: 'exec', labelKey: 'cap_exec', risk: 'high' }, + { bit: CAP_UPGRADE, key: 'upgrade', labelKey: 'cap_upgrade', risk: 'low' }, + { bit: CAP_PING_ICMP, key: 'ping_icmp', labelKey: 'cap_ping_icmp', risk: 'low' }, + { bit: CAP_PING_TCP, key: 'ping_tcp', labelKey: 'cap_ping_tcp', risk: 'low' }, + { bit: CAP_PING_HTTP, key: 'ping_http', labelKey: 'cap_ping_http', risk: 'low' }, + { bit: CAP_FILE, key: 'file', labelKey: 'cap_file', risk: 'high' }, + { bit: CAP_DOCKER, key: 'docker', labelKey: 'cap_docker', risk: 'high' }, + { bit: CAP_SECURITY_EVENTS, key: 'security_events', labelKey: 'cap_security_events', risk: 'low' }, + { bit: CAP_FIREWALL_BLOCK, key: 'firewall_block', labelKey: 'cap_firewall_block', risk: 'medium' }, + { bit: CAP_IP_QUALITY, key: 'ip_quality', labelKey: 'cap_ip_quality', risk: 'medium' } +] as const + +export type CapabilityRisk = (typeof CAPABILITIES)[number]['risk'] diff --git a/crates/common/examples/dump_capabilities_ts.rs b/crates/common/examples/dump_capabilities_ts.rs new file mode 100644 index 000000000..ad284aea3 --- /dev/null +++ b/crates/common/examples/dump_capabilities_ts.rs @@ -0,0 +1,57 @@ +//! Prints `apps/web/src/lib/capability-bits.generated.ts` from the +//! [`ALL_CAPABILITIES`] metadata table, so the web bit definitions cannot +//! drift from Rust. Regenerate via `bun run generate:capabilities` in +//! `apps/web` (same pipeline shape as `dump_openapi` → api-types). + +use serverbee_common::constants::{ALL_CAPABILITIES, CAP_DEFAULT}; + +fn main() { + let derived_default: u32 = ALL_CAPABILITIES + .iter() + .filter(|m| m.default_enabled) + .map(|m| m.bit) + .fold(0, |acc, bit| acc | bit); + assert_eq!( + derived_default, CAP_DEFAULT, + "CAP_DEFAULT must equal the OR of default_enabled bits" + ); + + println!("// AUTO-GENERATED from crates/common/src/constants.rs (ALL_CAPABILITIES)."); + println!("// Regenerate with `bun run generate:capabilities` — do not edit by hand."); + println!(); + for meta in ALL_CAPABILITIES { + println!( + "export const CAP_{} = {}", + meta.key.to_uppercase(), + meta.bit + ); + } + println!(); + let default_keys = ALL_CAPABILITIES + .iter() + .filter(|m| m.default_enabled) + .map(|m| m.key) + .collect::>() + .join(" + "); + println!("// The OR of every default_enabled bit: {default_keys}."); + println!("export const CAP_DEFAULT = {CAP_DEFAULT}"); + println!(); + println!("export const CAPABILITIES = ["); + let entries = ALL_CAPABILITIES + .iter() + .map(|meta| { + format!( + " {{ bit: CAP_{}, key: '{}', labelKey: 'cap_{}', risk: '{}' }}", + meta.key.to_uppercase(), + meta.key, + meta.key, + meta.risk_level + ) + }) + .collect::>() + .join(",\n"); + println!("{entries}"); + println!("] as const"); + println!(); + println!("export type CapabilityRisk = (typeof CAPABILITIES)[number]['risk']"); +} From 3acd704bba42343066b01c873bc0a8a5bde85a9f Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sun, 12 Jul 2026 01:35:54 +0800 Subject: [PATCH 38/44] docs: sync capability codegen note into AGENTS.md --- AGENTS.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/AGENTS.md b/AGENTS.md index 12f953667..a121dcf57 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -91,7 +91,7 @@ RBAC: Admin (full access) vs Member (read-only). `require_admin` middleware on w - **API responses**: All endpoints return `Json>` wrapping data in `{ data: T }` - **OpenAPI**: Every endpoint annotated with `#[utoipa::path]`, every DTO with `#[derive(ToSchema)]`. Swagger UI at `/swagger-ui/` - **Config**: Figment loads TOML then env vars. Prefix `SERVERBEE_`, nested separator `__` (double underscore). Example: `SERVERBEE_ADMIN__PASSWORD` → `admin.password`. **When adding/changing env vars, update `ENV.md` and `apps/docs/content/docs/{en,zh}/configuration.mdx` simultaneously.** -- **Capabilities**: u32 bitmask, defined in `crates/common/src/constants.rs` — `CAP_TERMINAL=1, CAP_EXEC=2, CAP_UPGRADE=4, CAP_PING_ICMP=8, CAP_PING_TCP=16, CAP_PING_HTTP=32, CAP_FILE=64, CAP_DOCKER=128, CAP_SECURITY_EVENTS=256, CAP_FIREWALL_BLOCK=512, CAP_IP_QUALITY=1024`. Default `CAP_DEFAULT=1852` (upgrade + ICMP/TCP/HTTP ping + security events + firewall blocklist + IP quality). **Capabilities are agent-owned**: the agent host computes its own bitmask from the `[capabilities]` config block (allow/deny over `CAP_DEFAULT`, see `compute_local_capabilities`) plus optional `--allow-cap`/`--deny-cap` CLI flags, and reports it in `SystemInfo.agent_local_capabilities`. The server **cannot** modify an agent's capabilities — it persists the reported value into `servers.capabilities` purely as a display mirror (`update_capabilities_mirror`) and gates control-plane requests on it. There is no server-side capability toggle; the web/iOS UI is read-only. +- **Capabilities**: u32 bitmask, defined in `crates/common/src/constants.rs` — `CAP_TERMINAL=1, CAP_EXEC=2, CAP_UPGRADE=4, CAP_PING_ICMP=8, CAP_PING_TCP=16, CAP_PING_HTTP=32, CAP_FILE=64, CAP_DOCKER=128, CAP_SECURITY_EVENTS=256, CAP_FIREWALL_BLOCK=512, CAP_IP_QUALITY=1024`. Default `CAP_DEFAULT=1852` (upgrade + ICMP/TCP/HTTP ping + security events + firewall blocklist + IP quality). **Capabilities are agent-owned**: the agent host computes its own bitmask from the `[capabilities]` config block (allow/deny over `CAP_DEFAULT`, see `compute_local_capabilities`) plus optional `--allow-cap`/`--deny-cap` CLI flags, and reports it in `SystemInfo.agent_local_capabilities`. The server **cannot** modify an agent's capabilities — it persists the reported value into `servers.capabilities` purely as a display mirror (`update_capabilities_mirror`) and gates control-plane requests on it. There is no server-side capability toggle; the web/iOS UI is read-only. Web bit definitions (`apps/web/src/lib/capability-bits.generated.ts`) are generated from `ALL_CAPABILITIES` — run `bun run generate:capabilities` in `apps/web` after changing capability metadata; iOS still mirrors by hand. - **Migrations**: sea-orm migrations in `crates/server/src/migration/`. Run automatically on startup. **Only implement `up()` — leave `down()` as a no-op (`Ok(())`).** Migrations are not reversible to avoid accidental data loss. ### Frontend From a3e8f6ee56db90f788c5a11aaef13a3f098bbd6a Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sun, 12 Jul 2026 01:45:34 +0800 Subject: [PATCH 39/44] refactor(server): name the IpChange transition helper The local closure-style helper in IpChange::is_transition was named `t`, which reads as opaque next to the semantic split between changed() and is_transition(). Rename it to `transitioned`. --- crates/server/src/router/ws/agent/system_info.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/crates/server/src/router/ws/agent/system_info.rs b/crates/server/src/router/ws/agent/system_info.rs index 96e33ae25..eaad815d1 100644 --- a/crates/server/src/router/ws/agent/system_info.rs +++ b/crates/server/src/router/ws/agent/system_info.rs @@ -61,12 +61,12 @@ impl IpChange { /// Drives the audit trail: first population happens on every fresh /// registration and must not spam the audit log. fn is_transition(&self) -> bool { - fn t(old: &Option, new: &Option) -> bool { + fn transitioned(old: &Option, new: &Option) -> bool { old.is_some() && old != new } - t(&self.old_ipv4, &self.new_ipv4) - || t(&self.old_ipv6, &self.new_ipv6) - || t(&self.old_remote_addr, &self.new_remote_addr) + transitioned(&self.old_ipv4, &self.new_ipv4) + || transitioned(&self.old_ipv6, &self.new_ipv6) + || transitioned(&self.old_remote_addr, &self.new_remote_addr) } fn detail(&self, server_id: &str) -> String { From fdd741bc3329b20944d3d3858b0fffba4a1676ba Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sun, 12 Jul 2026 01:45:34 +0800 Subject: [PATCH 40/44] test(server): pin IpChange change-vs-transition semantics The integration test only asserted that some ip_changed audit row exists, leaving the first-population rule unpinned. Add unit tests asserting that None -> Some drives changed() (alert + broadcast) but not is_transition() (audit), that Some -> different and Some -> None drive both, and that equal addresses drive neither. --- .../server/src/router/ws/agent/system_info.rs | 50 +++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/crates/server/src/router/ws/agent/system_info.rs b/crates/server/src/router/ws/agent/system_info.rs index eaad815d1..2adefb93a 100644 --- a/crates/server/src/router/ws/agent/system_info.rs +++ b/crates/server/src/router/ws/agent/system_info.rs @@ -437,3 +437,53 @@ async fn update_server_geo( active.update(db).await?; Ok(()) } + +#[cfg(test)] +mod tests { + use super::IpChange; + + fn change(old: Option<&str>, new: Option<&str>) -> IpChange { + IpChange { + old_ipv4: old.map(str::to_string), + new_ipv4: new.map(str::to_string), + old_ipv6: None, + new_ipv6: None, + old_remote_addr: None, + new_remote_addr: None, + } + } + + /// First population (None → Some) happens on every fresh registration: it + /// must drive the alert/broadcast path but never the audit trail. + #[test] + fn first_population_changes_without_being_a_transition() { + let c = change(None, Some("1.2.3.4")); + assert!(c.changed()); + assert!(!c.is_transition()); + } + + /// A real address move (Some → different Some) drives both paths. + #[test] + fn address_move_is_a_transition() { + let c = change(Some("1.2.3.4"), Some("5.6.7.8")); + assert!(c.changed()); + assert!(c.is_transition()); + } + + /// Losing an address (Some → None) is a transition too — it must be + /// audited, not mistaken for first population. + #[test] + fn address_loss_is_a_transition() { + let c = change(Some("1.2.3.4"), None); + assert!(c.changed()); + assert!(c.is_transition()); + } + + #[test] + fn identical_addresses_are_no_change_at_all() { + let c = change(Some("1.2.3.4"), Some("1.2.3.4")); + assert!(!c.changed()); + assert!(!c.is_transition()); + assert!(!change(None, None).changed()); + } +} From d10c83d460782e93c71d5400043f7e1f1b7856e7 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sun, 12 Jul 2026 01:45:47 +0800 Subject: [PATCH 41/44] fix(server): make non-alertable metric columns unresolvable in alerts rollup::alert_metric returned 0.0 for unknown rule types and for the deliberately non-alertable transfer counters, so a degenerate threshold (min <= 0) could still fire on them. Return Option instead and skip valueless samples in check_threshold, turning the documented "cannot be alerted on" claim into a hard guarantee. Behavior only changes for those degenerate configurations. Also correct the RollupAgg::MaxInt doc: MAX() keeps the window maximum, which equals the window-end value only while the counter grows monotonically. --- crates/server/src/service/alert.rs | 6 ++- crates/server/src/service/rollup.rs | 66 +++++++++++++++++------------ 2 files changed, 44 insertions(+), 28 deletions(-) diff --git a/crates/server/src/service/alert.rs b/crates/server/src/service/alert.rs index 2dafc4cbf..43ad50962 100644 --- a/crates/server/src/service/alert.rs +++ b/crates/server/src/service/alert.rs @@ -1031,7 +1031,11 @@ async fn check_threshold(db: &DatabaseConnection, server_id: &str, item: &AlertR let total = records.len(); for rec in &records { - let value = rollup::alert_metric(rec, &item.rule_type); + // Unresolvable rule types (unknown, or non-alertable transfer + // counters) contribute no sample, so such rules never reach 70%. + let Some(value) = rollup::alert_metric(rec, &item.rule_type) else { + continue; + }; let exceeds = match (item.min, item.max) { (Some(min), Some(max)) => value >= min && value <= max, (Some(min), None) => value >= min, diff --git a/crates/server/src/service/rollup.rs b/crates/server/src/service/rollup.rs index e49e99cdf..dbc9cddf3 100644 --- a/crates/server/src/service/rollup.rs +++ b/crates/server/src/service/rollup.rs @@ -48,8 +48,9 @@ pub enum RollupAgg { Avg, /// `CAST(AVG(col) AS INTEGER)` — integer gauges. AvgInt, - /// `CAST(MAX(col) AS INTEGER)` — cumulative counters keep the window-end - /// value instead of a meaningless average. + /// `CAST(MAX(col) AS INTEGER)` — cumulative counters keep the window + /// maximum (equal to the window-end value while the counter grows + /// monotonically) instead of a meaningless average. MaxInt, } @@ -178,13 +179,15 @@ pub const METRIC_COLUMNS: &[MetricColumn] = &[ }, ]; -/// Metric value an alert rule reads from a raw record. Unknown rule types -/// read 0.0 (the legacy fallback) so a misconfigured rule never fires. -pub fn alert_metric(rec: &record::Model, rule_type: &str) -> f64 { +/// Metric value an alert rule reads from a raw record. `None` for unknown +/// rule types and the deliberately non-alertable columns (`alert_rule_type: +/// None`), so such rules cannot fire — not even a `min <= 0` threshold that +/// a 0.0 fallback would satisfy. +pub fn alert_metric(rec: &record::Model, rule_type: &str) -> Option { METRIC_COLUMNS .iter() .find(|c| c.alert_rule_type == Some(rule_type)) - .map_or(0.0, |c| (c.read)(rec)) + .map(|c| (c.read)(rec)) } /// Build the hourly rollup upsert. Placeholders: bucket time, window start, @@ -326,11 +329,17 @@ mod tests { } } + #[track_caller] + fn assert_metric(rec: &record::Model, rule_type: &str, expected: f64) { + let value = alert_metric(rec, rule_type).expect("alertable rule type"); + assert!((value - expected).abs() < f64::EPSILON); + } + #[test] fn alert_metric_reads_direct_columns() { let rec = make_record(85.5, 4_000_000, 1.2); - assert!((alert_metric(&rec, "cpu") - 85.5).abs() < f64::EPSILON); - assert!((alert_metric(&rec, "load1") - 1.2).abs() < f64::EPSILON); + assert_metric(&rec, "cpu", 85.5); + assert_metric(&rec, "load1", 1.2); } /// Aliased rule types ("memory", "process", "gpu"…) resolve through the @@ -338,21 +347,23 @@ mod tests { #[test] fn alert_metric_resolves_rule_type_aliases() { let rec = make_record(50.0, 8_000_000, 0.0); - assert!((alert_metric(&rec, "memory") - 8_000_000.0).abs() < f64::EPSILON); - assert!((alert_metric(&rec, "process") - 200.0).abs() < f64::EPSILON); - assert!((alert_metric(&rec, "gpu") - 40.0).abs() < f64::EPSILON); - assert!((alert_metric(&rec, "tcp_conn") - 100.0).abs() < f64::EPSILON); - assert!((alert_metric(&rec, "udp_conn") - 50.0).abs() < f64::EPSILON); - assert!((alert_metric(&rec, "temperature") - 55.0).abs() < f64::EPSILON); + assert_metric(&rec, "memory", 8_000_000.0); + assert_metric(&rec, "process", 200.0); + assert_metric(&rec, "gpu", 40.0); + assert_metric(&rec, "tcp_conn", 100.0); + assert_metric(&rec, "udp_conn", 50.0); + assert_metric(&rec, "temperature", 55.0); } /// Unknown rule types and the deliberately non-alertable transfer - /// counters read 0.0 so such rules can never fire. + /// counters resolve to `None` so such rules can never fire — even a + /// degenerate `min <= 0` threshold gets no value to compare against. #[test] - fn alert_metric_unknown_and_non_alertable_read_zero() { + fn alert_metric_unknown_and_non_alertable_resolve_to_none() { let rec = make_record(99.0, 0, 0.0); - assert!((alert_metric(&rec, "nonexistent") - 0.0).abs() < f64::EPSILON); - assert!((alert_metric(&rec, "net_in_transfer") - 0.0).abs() < f64::EPSILON); + assert_eq!(alert_metric(&rec, "nonexistent"), None); + assert_eq!(alert_metric(&rec, "net_in_transfer"), None); + assert_eq!(alert_metric(&rec, "net_out_transfer"), None); } #[test] @@ -365,21 +376,22 @@ mod tests { rec.load15 = 3.5; rec.net_in_speed = 111; rec.net_out_speed = 222; - assert!((alert_metric(&rec, "swap") - 1024.0).abs() < f64::EPSILON); - assert!((alert_metric(&rec, "disk") - 2048.0).abs() < f64::EPSILON); - assert!((alert_metric(&rec, "load5") - 2.5).abs() < f64::EPSILON); - assert!((alert_metric(&rec, "load15") - 3.5).abs() < f64::EPSILON); - assert!((alert_metric(&rec, "net_in_speed") - 111.0).abs() < f64::EPSILON); - assert!((alert_metric(&rec, "net_out_speed") - 222.0).abs() < f64::EPSILON); + assert_metric(&rec, "swap", 1024.0); + assert_metric(&rec, "disk", 2048.0); + assert_metric(&rec, "load5", 2.5); + assert_metric(&rec, "load15", 3.5); + assert_metric(&rec, "net_in_speed", 111.0); + assert_metric(&rec, "net_out_speed", 222.0); } #[test] fn alert_metric_none_temperature_and_gpu_read_zero() { - // None temperature/gpu must coerce to 0.0 via unwrap_or. + // Absent sensors are an alertable column reading 0.0 (Some), not an + // unknown rule type (None). let mut rec = make_record(0.0, 0, 0.0); rec.temperature = None; rec.gpu_usage = None; - assert!((alert_metric(&rec, "temperature") - 0.0).abs() < f64::EPSILON); - assert!((alert_metric(&rec, "gpu") - 0.0).abs() < f64::EPSILON); + assert_metric(&rec, "temperature", 0.0); + assert_metric(&rec, "gpu", 0.0); } } From 013bfe8311905ad184e4709aeece70d1587dc3fe Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sun, 12 Jul 2026 01:45:47 +0800 Subject: [PATCH 42/44] docs: align rollup wording with window-maximum semantics ADR-0003 and CONTEXT.md described the transfer-counter rollup as "window-end MAX"; MAX() is the window maximum (coinciding with the window-end value only for monotonic counters). Also record that alert_metric now resolves non-alertable columns to None. --- CONTEXT.md | 2 +- .../0003-metric-columns-owned-by-rollup-descriptor.md | 11 +++++++---- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index cc2e792e1..0875b9d83 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -6,5 +6,5 @@ - **Standalone public network page** — `/status/network/$serverId`; exists only as a fallback for status pages configured with `show_server_detail=false` but `show_network=true` (see ADR-0001). - **Live metrics (实时指标帧)** — the partial projection of a server carried by WS `update` frames: the fields an agent report can populate (usage, speeds, loads, connection counts), plus `name` kept purely for decoder compatibility. Static facts are absent from the wire; clients keep their cached values when merging (see ADR-0002). _Avoid_: partial status, update payload. - **Full server status** — the complete ~35-field snapshot (`ServerStatus`) carried by `full_sync` and REST. The only sources allowed to seed the catalog or overwrite static facts (totals, os, tags, geo, enrollment). -- **Rollup policy (降采样策略)** — how raw metric records compress over time: per-column aggregation (AVG vs window-end MAX), the raw/hourly table switch point, and which columns alerts may read. Owned by `service::rollup`; declared per column in the **metric column descriptor** (see ADR-0003). _Avoid_: aggregation logic, downsampling code. +- **Rollup policy (降采样策略)** — how raw metric records compress over time: per-column aggregation (AVG vs window MAX), the raw/hourly table switch point, and which columns alerts may read. Owned by `service::rollup`; declared per column in the **metric column descriptor** (see ADR-0003). _Avoid_: aggregation logic, downsampling code. - **Metric column descriptor** — the one-row-per-column table (`METRIC_COLUMNS`) declaring a scalar metric's SQL name, rollup aggregation, alert rule type, and typed accessor. Adding a scalar metric means adding one descriptor row (plus entity/migration, which the compiler enforces). The web analogue is `METRIC_CHART_SPECS` for chart display. diff --git a/docs/adr/0003-metric-columns-owned-by-rollup-descriptor.md b/docs/adr/0003-metric-columns-owned-by-rollup-descriptor.md index 0433afad7..b9bc9293e 100644 --- a/docs/adr/0003-metric-columns-owned-by-rollup-descriptor.md +++ b/docs/adr/0003-metric-columns-owned-by-rollup-descriptor.md @@ -19,13 +19,16 @@ the public DTO maps) had no compiler protection. `service::rollup` owns the rollup policy. Every scalar column of `records`/`records_hourly` is declared once in `METRIC_COLUMNS` with its SQL -name, `RollupAgg` (AVG / CAST-AVG / window-end MAX), optional alert rule type, -and a typed accessor. From that table we derive: +name, `RollupAgg` (AVG / CAST-AVG / window MAX — for a monotonically growing +counter the maximum coincides with the window-end value), optional alert rule +type, and a typed accessor. From that table we derive: - the hourly rollup upsert (`aggregate_hourly_sql()`), pinned byte-for-byte to the original hand-written SQL by a golden test; -- the alert metric read (`alert_metric`), with transfer counters deliberately - non-alertable (`alert_rule_type: None`); +- the alert metric read (`alert_metric` → `Option`), with transfer + counters deliberately non-alertable: `alert_rule_type: None` resolves to + `None`, and `check_threshold` skips samples without a value, so no threshold + shape (including `min <= 0`) can fire on them; - the raw/hourly switch (`select_history_table`, `RAW_WINDOW_MAX_HOURS`). Alert evaluation reads through `RecordService::query_recent` / From b3281be1303be291491cbb615d7fd8f493e0b6c3 Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sun, 12 Jul 2026 01:45:47 +0800 Subject: [PATCH 43/44] docs(agent): declare collector elapsed-window stamping order Collector::collect stamps prev_time before assemble_report reads the counters, so the window is measured refresh-to-refresh. State that explicitly and note the us-level skew is absorbed by the 1s clamp. --- crates/agent/src/collector/mod.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/crates/agent/src/collector/mod.rs b/crates/agent/src/collector/mod.rs index 123067aae..d52efc0c9 100644 --- a/crates/agent/src/collector/mod.rs +++ b/crates/agent/src/collector/mod.rs @@ -55,6 +55,10 @@ impl Collector { pub fn collect(&mut self) -> SystemReport { self.source.refresh(); + // prev_time is stamped before the counters are read inside + // assemble_report, so the window is measured refresh-to-refresh + // rather than read-to-read. The µs-level skew this introduces is + // absorbed by the 1s elapsed clamp below. let elapsed = self.prev_time.elapsed().as_secs_f64(); self.prev_time = Instant::now(); self.assemble_report(elapsed) From c8ab56e0e2780dd663d7763a4f511d8b61e01feb Mon Sep 17 00:00:00 2001 From: ZingerLittleBee <6970999@gmail.com> Date: Sun, 12 Jul 2026 02:51:24 +0800 Subject: [PATCH 44/44] fix(agent): iterate network values directly to satisfy for_kv_map CI runs a newer clippy (1.97) that sees through sysinfo's Networks Deref to the underlying map and flags the (_name, data) iteration. Use networks.values() in total_bytes and its test. --- crates/agent/src/collector/network.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/agent/src/collector/network.rs b/crates/agent/src/collector/network.rs index d26e666b0..992ee0956 100644 --- a/crates/agent/src/collector/network.rs +++ b/crates/agent/src/collector/network.rs @@ -2,7 +2,7 @@ use sysinfo::Networks; pub fn total_bytes(networks: &Networks) -> (u64, u64) { let (mut total_in, mut total_out) = (0u64, 0u64); - for (_name, data) in networks.iter() { + for data in networks.values() { total_in += data.total_received(); total_out += data.total_transmitted(); } @@ -30,7 +30,7 @@ mod tests { // equals manual aggregation over the same snapshot. let mut sum_in = 0u64; let mut sum_out = 0u64; - for (_name, data) in networks.iter() { + for data in networks.values() { sum_in += data.total_received(); sum_out += data.total_transmitted(); }