diff --git a/wakey-control-plane/src/api/control/fleet/mod.rs b/wakey-control-plane/src/api/control/fleet/mod.rs index a0be3f9..502bf8e 100644 --- a/wakey-control-plane/src/api/control/fleet/mod.rs +++ b/wakey-control-plane/src/api/control/fleet/mod.rs @@ -4,17 +4,16 @@ use axum::Json; use axum::extract::{Query, State}; use axum::http::StatusCode; use axum::response::IntoResponse; -use futures_util::StreamExt; -use tracing::{Instrument, info, warn}; -use wakey_agent::protocol::{AgentCommand, InventoryRequest}; +use tracing::warn; use crate::api::ApiError; -use crate::api::commands::{RelayCommandResponse, relay_agent_command}; use crate::runtime::AppState; mod build; +mod refresh; mod types; mod wake; +pub use refresh::refresh_fleet_devices; pub use wake::wake_fleet_device; #[cfg(test)] @@ -24,10 +23,7 @@ use build::{ AgentRuntimeStatus, FleetBuildContext, build_fleet_devices, filter_fleet_devices, known_device_summary, }; -use types::{ - FleetDevice, ListFleetDevicesQuery, RefreshFleetAgentResult, RefreshFleetDevicesRequest, - RefreshFleetDevicesResponse, -}; +use types::{FleetDevice, ListFleetDevicesQuery}; pub(crate) type FleetCache = crate::runtime::view_cache::ViewCache<(u64, BTreeSet), Vec>; @@ -80,8 +76,6 @@ pub(crate) async fn sample_fleet_counts( Err(crate::runtime::view_cache::RebuildLimit.into()) } -const REFRESH_CONCURRENCY: usize = 8; - fn fleet_error(err: anyhow::Error, fallback_code: &str) -> ApiError { if err.is::() { return ApiError::new( @@ -111,146 +105,6 @@ pub async fn list_fleet_devices( } } -#[tracing::instrument(name = "fleet_refresh", skip_all, fields(agents = tracing::field::Empty, succeeded = tracing::field::Empty, failed = tracing::field::Empty, accepted = tracing::field::Empty))] -pub async fn refresh_fleet_devices( - State(state): State, - Json(req): Json, -) -> Result { - let agent_ids = if req.agent_ids.is_empty() { - let sessions = state.sessions.read().await; - sessions.keys().cloned().collect::>() - } else { - req.agent_ids.clone() - }; - let agent_ids = normalize_refresh_agent_ids(agent_ids); - tracing::Span::current().record("agents", agent_ids.len()); - - let results = futures_util::stream::iter(agent_ids.into_iter().map(|agent_id| { - let state = state.clone(); - // Each concurrent future carries its own span through relay and storage. - let span = tracing::info_span!("fleet_refresh_agent", agent_id = %agent_id, outcome = tracing::field::Empty, accepted = tracing::field::Empty); - async move { - let response = relay_agent_command( - &state, - &agent_id, - AgentCommand::Inventory(InventoryRequest { - query: None, - name: None, - ips: Vec::new(), - devs: Vec::new(), - nuds: Vec::new(), - macs: Vec::new(), - }), - req.timeout_ms, - ) - .await; - let result = refresh_fleet_agent_result(&state, &agent_id, response).await; - tracing::Span::current().record("outcome", result.status.as_str()); - tracing::Span::current().record("accepted", result.accepted); - if result.status != "ok" { - // The response retains diagnostic detail; do not export payloads. - warn!("fleet agent refresh failed"); - } - result - } - .instrument(span) - })) - .buffer_unordered(REFRESH_CONCURRENCY) - .collect::>() - .await; - - let mut results = results; - results.sort_by(|a, b| a.agent_id.cmp(&b.agent_id)); - let total_accepted = results.iter().map(|result| result.accepted).sum::(); - let succeeded = results - .iter() - .filter(|result| result.status == "ok") - .count(); - let span = tracing::Span::current(); - span.record("succeeded", succeeded); - span.record("failed", results.len() - succeeded); - span.record("accepted", total_accepted); - info!("fleet refresh completed"); - - Ok(( - StatusCode::OK, - Json(RefreshFleetDevicesResponse { - total_accepted, - agents: results, - }), - )) -} - -fn normalize_refresh_agent_ids(mut agent_ids: Vec) -> Vec { - agent_ids.sort(); - agent_ids.dedup(); - agent_ids -} - -async fn refresh_fleet_agent_result( - state: &AppState, - agent_id: &str, - response: Result, -) -> RefreshFleetAgentResult { - match response { - Ok(response) if response.status == "ok" => { - let Some(result) = response.result else { - return RefreshFleetAgentResult { - agent_id: agent_id.to_string(), - status: "error".into(), - accepted: 0, - error: Some("inventory command returned no result".into()), - }; - }; - match serde_json::from_value::(result) - .map(|inventory| inventory.devices) - { - Ok(devices) => { - match state - .store - .replace_agent_device_snapshot(agent_id, &devices) - .await - { - Ok(accepted) => RefreshFleetAgentResult { - agent_id: agent_id.to_string(), - status: "ok".into(), - accepted, - error: None, - }, - Err(err) => RefreshFleetAgentResult { - agent_id: agent_id.to_string(), - status: "error".into(), - accepted: 0, - error: Some(err.to_string()), - }, - } - } - Err(err) => RefreshFleetAgentResult { - agent_id: agent_id.to_string(), - status: "error".into(), - accepted: 0, - error: Some(err.to_string()), - }, - } - } - Ok(response) => RefreshFleetAgentResult { - agent_id: agent_id.to_string(), - status: "error".into(), - accepted: 0, - error: response - .error - .map(|error| error.message) - .or_else(|| Some("inventory command failed".into())), - }, - Err(err) => RefreshFleetAgentResult { - agent_id: agent_id.to_string(), - status: "error".into(), - accepted: 0, - error: Some(format!("{}: {}", err.status, err.message)), - }, - } -} - #[tracing::instrument(name = "fleet_load", skip_all, fields(grouped = tracing::field::Empty, filtered = tracing::field::Empty, returned = tracing::field::Empty))] async fn load_fleet_devices( state: &AppState, diff --git a/wakey-control-plane/src/api/control/fleet/refresh.rs b/wakey-control-plane/src/api/control/fleet/refresh.rs new file mode 100644 index 0000000..7ae69de --- /dev/null +++ b/wakey-control-plane/src/api/control/fleet/refresh.rs @@ -0,0 +1,337 @@ +//! Refresh agents independently: one failed relay or snapshot does not cancel others. + +use axum::{Json, extract::State, http::StatusCode, response::IntoResponse}; +use futures_util::StreamExt; +use tracing::{info, warn}; +use wakey_agent::protocol::{AgentCommand, InventoryRequest}; + +use super::types::{ + RefreshFleetAgentResult, RefreshFleetDevicesRequest, RefreshFleetDevicesResponse, +}; +use crate::api::{ + ApiError, + commands::{RelayCommandResponse, relay_agent_command}, +}; +use crate::runtime::AppState; +use crate::state::Store; + +const REFRESH_CONCURRENCY: usize = 8; + +#[tracing::instrument(name = "fleet_refresh", skip_all, fields(agents = tracing::field::Empty, succeeded = tracing::field::Empty, failed = tracing::field::Empty, accepted = tracing::field::Empty))] +pub async fn refresh_fleet_devices( + State(state): State, + Json(req): Json, +) -> Result { + let agent_ids = select_agents(&state, req.agent_ids).await; + tracing::Span::current().record("agents", agent_ids.len()); + let results = futures_util::stream::iter( + agent_ids + .into_iter() + .map(|agent_id| refresh_agent(&state, agent_id, req.timeout_ms)), + ) + .buffer_unordered(REFRESH_CONCURRENCY) + .collect::>() + .await; + let response = summarize(results); + let succeeded = response + .agents + .iter() + .filter(|result| result.status == "ok") + .count(); + let span = tracing::Span::current(); + span.record("succeeded", succeeded); + span.record("failed", response.agents.len() - succeeded); + span.record("accepted", response.total_accepted); + info!("fleet refresh completed"); + Ok((StatusCode::OK, Json(response))) +} + +/// Default to a point-in-time connected set. Explicit IDs remain eligible so +/// the relay can report disconnected agents as per-agent errors, not omit them. +async fn select_agents(state: &AppState, requested: Vec) -> Vec { + let ids = if requested.is_empty() { + state.sessions.read().await.keys().cloned().collect() + } else { + requested + }; + normalize_refresh_agent_ids(ids) +} + +pub(super) fn normalize_refresh_agent_ids(mut agent_ids: Vec) -> Vec { + agent_ids.sort(); + agent_ids.dedup(); + agent_ids +} + +#[tracing::instrument(name = "fleet_refresh_agent", skip_all, fields(agent_id = %agent_id, outcome = tracing::field::Empty, accepted = tracing::field::Empty))] +async fn refresh_agent( + state: &AppState, + agent_id: String, + timeout_ms: Option, +) -> RefreshFleetAgentResult { + let response = relay_agent_command( + state, + &agent_id, + AgentCommand::Inventory(InventoryRequest { + query: None, + name: None, + ips: Vec::new(), + devs: Vec::new(), + nuds: Vec::new(), + macs: Vec::new(), + }), + timeout_ms, + ) + .await; + let result = store_inventory(&state.store, &agent_id, response).await; + tracing::Span::current().record("outcome", result.status.as_str()); + tracing::Span::current().record("accepted", result.accepted); + if result.status != "ok" { + // The response retains diagnostic detail; do not export payloads. + warn!("fleet agent refresh failed"); + } + result +} + +/// Validate the relay result before replacing anything. An empty *valid* +/// inventory is still a complete snapshot and intentionally clears old rows. +fn decode_inventory( + response: Result, +) -> Result { + let response = response.map_err(|err| format!("{}: {}", err.status, err.message))?; + if response.status != "ok" { + return Err(response + .error + .map(|error| error.message) + .unwrap_or_else(|| "inventory command failed".into())); + } + let result = response + .result + .ok_or_else(|| "inventory command returned no result".to_string())?; + serde_json::from_value(result).map_err(|err| err.to_string()) +} + +async fn store_inventory( + store: &Store, + agent_id: &str, + response: Result, +) -> RefreshFleetAgentResult { + let accepted = async { + let inventory = decode_inventory(response)?; + store + .replace_agent_device_snapshot(agent_id, &inventory.devices) + .await + .map_err(|err| err.to_string()) + } + .await; + match accepted { + Ok(accepted) => RefreshFleetAgentResult { + agent_id: agent_id.into(), + status: "ok".into(), + accepted, + error: None, + }, + Err(error) => RefreshFleetAgentResult { + agent_id: agent_id.into(), + status: "error".into(), + accepted: 0, + error: Some(error), + }, + } +} + +/// Concurrent completion order is deliberately not part of the response contract. +fn summarize(mut agents: Vec) -> RefreshFleetDevicesResponse { + agents.sort_by(|a, b| a.agent_id.cmp(&b.agent_id)); + RefreshFleetDevicesResponse { + total_accepted: agents.iter().map(|result| result.accepted).sum(), + agents, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn response(result: Option) -> RelayCommandResponse { + RelayCommandResponse { + request_id: "refresh-test".into(), + status: "ok".into(), + result, + error: None, + } + } + + #[test] + fn decode_keeps_relay_and_payload_failures_distinct() { + assert_eq!( + decode_inventory(Ok(response(None))).unwrap_err(), + "inventory command returned no result" + ); + assert!(decode_inventory(Ok(response(Some(json!({"devices": "invalid"}))))).is_err()); + let mut failed = response(Some(json!({"devices": []}))); + failed.status = "error".into(); + assert_eq!( + decode_inventory(Ok(failed)).unwrap_err(), + "inventory command failed" + ); + assert_eq!( + decode_inventory(Err(ApiError::new( + StatusCode::NOT_FOUND, + "agent_not_connected", + "agent is not connected" + ))) + .unwrap_err(), + "404 Not Found: agent is not connected" + ); + assert!( + decode_inventory(Ok(response(Some(json!({"devices": []}))))) + .unwrap() + .devices + .is_empty() + ); + } + + #[test] + fn summary_keeps_partial_failures_and_sorts_by_agent() { + let result = summarize(vec![ + RefreshFleetAgentResult { + agent_id: "z".into(), + status: "ok".into(), + accepted: 2, + error: None, + }, + RefreshFleetAgentResult { + agent_id: "b".into(), + status: "error".into(), + accepted: 0, + error: Some("disconnected".into()), + }, + RefreshFleetAgentResult { + agent_id: "a".into(), + status: "ok".into(), + accepted: 3, + error: None, + }, + ]); + assert_eq!(result.total_accepted, 5); + assert_eq!( + result + .agents + .iter() + .map(|row| row.agent_id.as_str()) + .collect::>(), + ["a", "b", "z"] + ); + assert_eq!(result.agents[1].error.as_deref(), Some("disconnected")); + assert!(summarize(vec![]).agents.is_empty()); + } + + #[tokio::test] + async fn handler_reports_explicit_disconnected_agents_without_omitting_them() { + use std::{collections::HashMap, sync::Arc, time::Duration}; + use tokio::sync::{Mutex, RwLock}; + + let dir = + std::env::temp_dir().join(format!("wakey-refresh-handler-{}", uuid::Uuid::new_v4())); + let state = AppState { + store: Arc::new( + Store::load_or_init(&dir.join("state.sqlite3"), vec![], Duration::from_secs(60)) + .await + .unwrap(), + ), + sessions: Arc::new(RwLock::new(HashMap::new())), + pending: Arc::new(Mutex::new(HashMap::new())), + public_url: "http://localhost".into(), + command_timeout: Duration::from_secs(30), + enroll_token_ttl: Duration::from_secs(60), + terminals: crate::runtime::terminals::TerminalRegistry::new(), + alerts: crate::api::AlertsHub::new(), + fleet: Arc::new(super::super::FleetCache::default()), + }; + assert!(select_agents(&state, vec![]).await.is_empty()); + let response = refresh_fleet_devices( + State(state.clone()), + Json(RefreshFleetDevicesRequest { + agent_ids: vec!["missing-b".into(), "missing-a".into(), "missing-b".into()], + timeout_ms: None, + }), + ) + .await + .unwrap() + .into_response(); + assert_eq!(response.status(), StatusCode::OK); + let body = axum::body::to_bytes(response.into_body(), 4096) + .await + .unwrap(); + let result: RefreshFleetDevicesResponse = serde_json::from_slice(&body).unwrap(); + assert_eq!(result.total_accepted, 0); + assert_eq!(result.agents.len(), 2); + assert_eq!(result.agents[0].agent_id, "missing-a"); + assert_eq!(result.agents[1].agent_id, "missing-b"); + for agent in result.agents { + assert_eq!(agent.status, "error"); + assert_eq!(agent.accepted, 0); + assert_eq!( + agent.error.as_deref(), + Some("404 Not Found: agent is not connected") + ); + } + assert!(state.pending.lock().await.is_empty()); + drop(state); + std::fs::remove_dir_all(dir).unwrap(); + } + + #[tokio::test] + async fn refresh_stores_snapshots_but_invalid_results_leave_them_untouched() { + let dir = std::env::temp_dir().join(format!("wakey-refresh-{}", uuid::Uuid::new_v4())); + let store = Store::load_or_init( + &dir.join("state.sqlite3"), + vec![], + std::time::Duration::from_secs(60), + ) + .await + .unwrap(); + let device = wakey_core::Device::from_parts( + vec![wakey_core::NeighborEntry { + ip: "192.168.1.10".parse().unwrap(), + mac: Some("aa:bb:cc:dd:ee:01".parse().unwrap()), + dev: Some("br-lan".into()), + state: wakey_core::NeighborState::Reachable, + }], + vec![], + ); + let inventory = serde_json::to_value(wakey_core::DeviceInventory { + devices: vec![device], + }) + .unwrap(); + let stored = store_inventory(&store, "agent", Ok(response(Some(inventory)))).await; + assert_eq!(stored.status, "ok"); + assert_eq!(stored.accepted, 1); + let revision = store.fleet_revision(); + let invalid = store_inventory(&store, "agent", Ok(response(None))).await; + assert_eq!(invalid.status, "error"); + assert_eq!(invalid.accepted, 0); + assert_eq!(store.fleet_revision(), revision); + let rows = store + .list_agent_device_rows_for_agent("agent") + .await + .unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].endpoints.len(), 1); + let empty = + store_inventory(&store, "agent", Ok(response(Some(json!({"devices": []}))))).await; + assert_eq!(empty.status, "ok"); + assert_eq!(empty.accepted, 0); + assert!( + store + .list_agent_device_rows_for_agent("agent") + .await + .unwrap() + .is_empty() + ); + drop(store); + std::fs::remove_dir_all(dir).unwrap(); + } +} diff --git a/wakey-control-plane/src/api/control/fleet/tests.rs b/wakey-control-plane/src/api/control/fleet/tests.rs index d703d75..8265616 100644 --- a/wakey-control-plane/src/api/control/fleet/tests.rs +++ b/wakey-control-plane/src/api/control/fleet/tests.rs @@ -8,7 +8,7 @@ use super::build::{ AgentRuntimeStatus, FleetBuildContext, build_fleet_devices, filter_fleet_devices, known_device_summary, }; -use super::normalize_refresh_agent_ids; +use super::refresh::normalize_refresh_agent_ids; use super::types::ListFleetDevicesQuery; use crate::state::{AgentDeviceRow, AgentDeviceWithChildren, DeviceIdentifier, KnownDevice};