diff --git a/ui/src/pages/TerminalPage.tsx b/ui/src/pages/TerminalPage.tsx index 29182c2..0ee4427 100644 --- a/ui/src/pages/TerminalPage.tsx +++ b/ui/src/pages/TerminalPage.tsx @@ -576,7 +576,10 @@ export function TerminalPage({ }); resizeObserver.observe(hostRef.current); void restoreTerminalSession(); + let refreshInFlight = false; const refreshSessions = () => { + if (refreshInFlight) return; + refreshInFlight = true; void listTerminals() .then((listed) => { if (!cancelled) { @@ -588,11 +591,13 @@ export function TerminalPage({ .catch(() => { // The attached terminal transport remains authoritative while a // background list refresh is temporarily unavailable. + }) + .finally(() => { + refreshInFlight = false; }); }; refreshSessionsRef.current = refreshSessions; - // Session lifecycle only changes at connect/disconnect moments or when the - // operator returns to the tab; refresh on those instead of polling. + const refreshTimer = window.setInterval(refreshSessions, 15_000); window.addEventListener("focus", refreshSessions); return () => { @@ -600,6 +605,7 @@ export function TerminalPage({ input.dispose(); titleChange.dispose(); resizeObserver.disconnect(); + window.clearInterval(refreshTimer); window.removeEventListener("focus", refreshSessions); refreshSessionsRef.current = () => {}; window.cancelAnimationFrame(resizeFrame); diff --git a/wakey-agent/src/terminal.rs b/wakey-agent/src/terminal.rs index df2802d..619c24b 100644 --- a/wakey-agent/src/terminal.rs +++ b/wakey-agent/src/terminal.rs @@ -522,7 +522,7 @@ async fn run_terminal_relay(relay: RelayConnection) -> Result<()> { let (stream, _) = tokio_tungstenite::connect_async(ws_url.as_str()) .await .context("failed to connect terminal relay websocket")?; - let (mut sink, mut source) = stream.split(); + let (mut sink, source) = stream.split(); send_json( &mut sink, &TerminalAgentHandshake::Auth { @@ -541,40 +541,61 @@ async fn run_terminal_relay(relay: RelayConnection) -> Result<()> { .await .map_err(|_| anyhow::anyhow!("terminal worker stopped"))?; - let mut output = relay.output_rx; + relay_loop(source, sink, relay.input, relay.generation, relay.output_rx).await +} + +async fn relay_loop( + mut source: St, + mut sink: S, + input: mpsc::Sender, + generation: u64, + mut output: mpsc::Receiver, +) -> Result<()> +where + S: futures_util::Sink + Unpin, + St: futures_util::Stream< + Item = std::result::Result, + > + Unpin, +{ + let mut worker_stopped = false; + let mut worker_stop_deadline: Option = None; + // Reserve at most one control so a full worker queue cannot block output. + let mut pending_input = None; + let mut close_after_input = false; loop { tokio::select! { biased; - incoming = source.next() => { + _ = async { + match worker_stop_deadline { + Some(deadline) => tokio::time::sleep_until(deadline).await, + None => std::future::pending().await, + } + } => break, + incoming = source.next(), if !worker_stopped && pending_input.is_none() => { let Some(message) = incoming else { break; }; match message.context("terminal relay websocket receive failed")? { Message::Binary(bytes) => { - relay.input.send(RelayInput::Binary { - generation: relay.generation, + pending_input = Some(RelayInput::Binary { + generation, bytes: bytes.to_vec(), - }).await.map_err(|_| anyhow::anyhow!("terminal worker stopped"))?; + }); } Message::Text(text) => match serde_json::from_str::(&text) .context("invalid terminal control frame")? { TerminalControl::Resize { rows, cols } => { - relay.input.send(RelayInput::Resize { - generation: relay.generation, + pending_input = Some(RelayInput::Resize { + generation, rows, cols, - }).await - .map_err(|_| anyhow::anyhow!("terminal worker stopped"))?; + }); } TerminalControl::Snapshot => { - relay.input.send(RelayInput::Snapshot { - generation: relay.generation, - }).await.map_err(|_| anyhow::anyhow!("terminal worker stopped"))?; + pending_input = Some(RelayInput::Snapshot { generation }); } TerminalControl::Close => { - let _ = relay.input.send(RelayInput::Close { - generation: relay.generation, - }).await; - break; + pending_input = Some(RelayInput::Close { generation }); + close_after_input = true; } _ => anyhow::bail!("terminal control frame has invalid direction"), }, @@ -586,9 +607,32 @@ async fn run_terminal_relay(relay: RelayConnection) -> Result<()> { Message::Frame(_) => {} } } + permit = input.reserve(), if !worker_stopped && pending_input.is_some() => { + match permit { + Ok(permit) => { + permit.send(pending_input.take().expect("pending terminal input")); + if close_after_input { + break; + } + } + Err(_) => { + worker_stopped = true; + pending_input = None; + worker_stop_deadline = + Some(tokio::time::Instant::now() + PTY_EXIT_DRAIN_TIMEOUT); + } + } + } outgoing = output.recv() => { let Some(message) = outgoing else { break; }; - sink.send(message).await.context("failed to send terminal relay output")?; + if let Some(deadline) = worker_stop_deadline { + tokio::time::timeout_at(deadline, sink.send(message)) + .await + .context("terminal relay output drain timed out")? + .context("failed to send terminal relay output")?; + } else { + sink.send(message).await.context("failed to send terminal relay output")?; + } } } } @@ -871,6 +915,58 @@ mod tests { } } + #[tokio::test] + async fn relay_drains_final_output_after_worker_stops_during_concurrent_controls() { + let (input, input_rx) = mpsc::channel(RELAY_INPUT_QUEUE); + drop(input_rx); + let (output_tx, output_rx) = mpsc::channel(RELAY_OUTPUT_QUEUE); + output_tx + .send(Message::Binary(vec![1, 2, 3].into())) + .await + .expect("live output queues"); + let exit_frame = serde_json::to_string(&TerminalControl::Exited { exit_code: Some(0) }) + .expect("exit frame encodes"); + output_tx + .send(Message::Text(exit_frame.clone().into())) + .await + .expect("exit output queues"); + drop(output_tx); + + let source = futures_util::stream::iter([ + Ok(Message::Binary(vec![4].into())), + Ok(Message::Text( + serde_json::to_string(&TerminalControl::Resize { + rows: 30, + cols: 100, + }) + .expect("resize encodes") + .into(), + )), + Ok(Message::Text( + serde_json::to_string(&TerminalControl::Snapshot) + .expect("snapshot encodes") + .into(), + )), + ]); + let delivered = Arc::new(Mutex::new(Vec::new())); + let sink = futures_util::sink::unfold(delivered.clone(), |delivered, frame| async move { + delivered.lock().expect("delivered frames").push(frame); + Ok::<_, tokio_tungstenite::tungstenite::Error>(delivered) + }); + + relay_loop(source, Box::pin(sink), input, 7, output_rx) + .await + .expect("relay drains after worker completion"); + + let delivered = delivered.lock().expect("delivered frames"); + assert_eq!(delivered.len(), 2); + assert_eq!(delivered[0].clone().into_data(), vec![1, 2, 3]); + let Message::Text(exit) = &delivered[1] else { + panic!("exit control should follow final bytes"); + }; + assert_eq!(exit.as_str(), exit_frame); + } + #[test] fn open_rejects_invalid_initial_size_before_spawning() { let mut config = crate::config::DEFAULT_CONFIG.clone(); diff --git a/wakey-control-plane/src/api/alerts.rs b/wakey-control-plane/src/api/alerts.rs index cd04d7d..4a2c6ac 100644 --- a/wakey-control-plane/src/api/alerts.rs +++ b/wakey-control-plane/src/api/alerts.rs @@ -35,6 +35,7 @@ pub struct AlertHistoryQuery { } #[derive(Debug, Clone, Deserialize)] +#[serde(default)] pub struct AlertRuleConfig { pub lookback_seconds: Option, pub timeout_threshold: Option, @@ -53,6 +54,45 @@ impl Default for AlertRuleConfig { } } +impl AlertRuleConfig { + fn validate(&self) -> Result<(), ApiError> { + let defaults = Self::default(); + let overrides = [ + ( + "lookback_seconds", + self.lookback_seconds, + defaults.lookback_seconds, + ), + ( + "timeout_threshold", + self.timeout_threshold, + defaults.timeout_threshold, + ), + ( + "auth_rejected_threshold", + self.auth_rejected_threshold, + defaults.auth_rejected_threshold, + ), + ( + "enroll_rejected_threshold", + self.enroll_rejected_threshold, + defaults.enroll_rejected_threshold, + ), + ]; + if let Some((field, _, _)) = overrides + .into_iter() + .find(|(_, requested, default)| requested.is_some() && requested != default) + { + return Err(ApiError::new( + StatusCode::BAD_REQUEST, + "unsupported_alert_override", + format!("custom alert {field} is not supported; omit it or use its default value"), + )); + } + Ok(()) + } +} + #[derive(Clone, Default)] pub struct AlertsSnapshot { pub ts_unix: u64, @@ -67,8 +107,17 @@ pub struct AlertsHub { tx: broadcast::Sender, latest_payload: Arc>>, latest_alerts: Arc>>, + evaluator_status: Arc>, } +#[derive(Default)] +struct EvaluatorStatus { + last_success_unix: Option, + failure: Option, +} + +const ALERT_SNAPSHOT_MAX_AGE_SECS: u64 = 15; + impl Default for AlertsHub { fn default() -> Self { Self::new() @@ -82,6 +131,7 @@ impl AlertsHub { tx, latest_payload: Arc::new(RwLock::new(None)), latest_alerts: Arc::new(RwLock::new(Vec::new())), + evaluator_status: Arc::new(RwLock::new(EvaluatorStatus::default())), } } @@ -97,6 +147,22 @@ impl AlertsHub { self.latest_payload.read().await.clone() } + async fn mark_unavailable(&self, reason: impl Into) { + self.evaluator_status.write().await.failure = Some(reason.into()); + } + + async fn unavailable_reason(&self) -> Option { + let status = self.evaluator_status.read().await; + if let Some(failure) = &status.failure { + return Some(failure.clone()); + } + match status.last_success_unix { + Some(ts) if now_unix().saturating_sub(ts) <= ALERT_SNAPSHOT_MAX_AGE_SECS => None, + Some(_) => Some("alert evaluator snapshot is stale".into()), + None => Some("alert evaluator has not produced a snapshot yet".into()), + } + } + pub async fn publish(&self, snapshot: AlertsSnapshot) { let payload = serde_json::json!({ "type": "alerts_snapshot", @@ -110,6 +176,10 @@ impl AlertsHub { }; *self.latest_alerts.write().await = snapshot.alerts; *self.latest_payload.write().await = Some(encoded.clone()); + let mut status = self.evaluator_status.write().await; + status.last_success_unix = Some(snapshot.ts_unix); + status.failure = None; + drop(status); let _ = self.tx.send(encoded); } } @@ -127,6 +197,7 @@ pub async fn run_alert_evaluator(state: AppState) { Ok(alerts) => alerts, Err(err) => { warn!(code = %err.code, "failed to evaluate alerts in background task"); + state.alerts.mark_unavailable(err.message.clone()).await; continue; } }; @@ -154,7 +225,18 @@ pub async fn run_alert_evaluator(state: AppState) { } } -pub async fn active_alerts(State(state): State) -> Result { +pub async fn active_alerts( + State(state): State, + Query(config): Query, +) -> Result { + config.validate()?; + if let Some(reason) = state.alerts.unavailable_reason().await { + return Err(ApiError::new( + StatusCode::SERVICE_UNAVAILABLE, + "alerts_unavailable", + reason, + )); + } let alerts = state.alerts.latest_alerts().await; Ok((StatusCode::OK, Json(alerts))) } @@ -426,9 +508,11 @@ fn now_unix() -> u64 { mod tests { use std::collections::HashSet; + use axum::http::StatusCode; + use crate::state::AuditEvent; - use super::{AlertBuildContext, build_alerts}; + use super::{AlertBuildContext, AlertRuleConfig, AlertsHub, AlertsSnapshot, build_alerts}; fn event(agent_id: Option<&str>, event_type: &str, outcome: &str, ts_unix: u64) -> AuditEvent { AuditEvent { @@ -517,4 +601,81 @@ mod tests { .any(|a| a.kind == "enroll_reject_spike" && a.agent_id.is_none()) ); } + + #[test] + fn alert_rule_config_accepts_omitted_and_default_values() { + AlertRuleConfig::default().validate().unwrap(); + AlertRuleConfig { + lookback_seconds: None, + timeout_threshold: None, + auth_rejected_threshold: None, + enroll_rejected_threshold: None, + } + .validate() + .unwrap(); + } + + #[test] + fn alert_rule_config_rejects_custom_overrides() { + let configs = [ + AlertRuleConfig { + lookback_seconds: Some(300), + ..Default::default() + }, + AlertRuleConfig { + timeout_threshold: Some(4), + ..Default::default() + }, + AlertRuleConfig { + auth_rejected_threshold: Some(4), + ..Default::default() + }, + AlertRuleConfig { + enroll_rejected_threshold: Some(6), + ..Default::default() + }, + ]; + + for config in configs { + let error = config.validate().unwrap_err(); + assert_eq!(error.status, StatusCode::BAD_REQUEST); + assert_eq!(error.code, "unsupported_alert_override"); + assert!(error.message.contains("not supported")); + } + } + + #[tokio::test] + async fn alert_hub_reports_initial_failure_and_recovers_on_snapshot() { + let hub = AlertsHub::new(); + assert!(hub.unavailable_reason().await.is_some()); + + hub.mark_unavailable("store query failed").await; + assert_eq!( + hub.unavailable_reason().await.as_deref(), + Some("store query failed") + ); + + hub.publish(AlertsSnapshot { + ts_unix: super::now_unix(), + alerts: Vec::new(), + recent_transitions: Vec::new(), + }) + .await; + assert_eq!(hub.unavailable_reason().await, None); + } + + #[tokio::test] + async fn alert_hub_rejects_stale_snapshots() { + let hub = AlertsHub::new(); + hub.publish(AlertsSnapshot { + ts_unix: 0, + alerts: Vec::new(), + recent_transitions: Vec::new(), + }) + .await; + assert_eq!( + hub.unavailable_reason().await.as_deref(), + Some("alert evaluator snapshot is stale") + ); + } } diff --git a/wakey-control-plane/src/api/control/fleet/mod.rs b/wakey-control-plane/src/api/control/fleet/mod.rs index e0bace1..fecfc8c 100644 --- a/wakey-control-plane/src/api/control/fleet/mod.rs +++ b/wakey-control-plane/src/api/control/fleet/mod.rs @@ -27,7 +27,6 @@ use types::{ RefreshFleetDevicesResponse, WakeFleetDeviceRequest, WakeFleetDeviceResponse, }; -const MAX_REFRESH_AGENTS: usize = 128; const REFRESH_CONCURRENCY: usize = 8; pub async fn list_fleet_devices( @@ -51,22 +50,13 @@ pub async fn refresh_fleet_devices( State(state): State, Json(req): Json, ) -> Result { - let mut agent_ids = if req.agent_ids.is_empty() { + let agent_ids = if req.agent_ids.is_empty() { let sessions = state.sessions.read().await; sessions.keys().cloned().collect::>() } else { req.agent_ids.clone() }; - agent_ids.sort(); - agent_ids.dedup(); - if agent_ids.len() > MAX_REFRESH_AGENTS { - warn!( - requested = agent_ids.len(), - capped = MAX_REFRESH_AGENTS, - "capping fleet refresh agent count" - ); - agent_ids.truncate(MAX_REFRESH_AGENTS); - } + let agent_ids = normalize_refresh_agent_ids(agent_ids); let results = futures_util::stream::iter(agent_ids.into_iter().map(|agent_id| { let state = state.clone(); @@ -105,6 +95,12 @@ pub async fn refresh_fleet_devices( )) } +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, diff --git a/wakey-control-plane/src/api/control/fleet/tests.rs b/wakey-control-plane/src/api/control/fleet/tests.rs index a611de8..c3dd62e 100644 --- a/wakey-control-plane/src/api/control/fleet/tests.rs +++ b/wakey-control-plane/src/api/control/fleet/tests.rs @@ -8,9 +8,23 @@ use super::build::{ AgentRuntimeStatus, FleetBuildContext, build_fleet_devices, filter_fleet_devices, known_device_summary, }; +use super::normalize_refresh_agent_ids; use super::types::ListFleetDevicesQuery; use crate::state::{AgentDeviceRow, AgentDeviceWithChildren, DeviceIdentifier, KnownDevice}; +#[test] +fn refresh_agent_selection_keeps_all_agents_beyond_previous_cap() { + let ids = (0..140) + .map(|index| format!("agent-{index:03}")) + .collect::>(); + + let selected = normalize_refresh_agent_ids(ids); + + assert_eq!(selected.len(), 140); + assert_eq!(selected.first().map(String::as_str), Some("agent-000")); + assert_eq!(selected.last().map(String::as_str), Some("agent-139")); +} + fn context(connected: &[&str]) -> FleetBuildContext { FleetBuildContext { agent_status: connected