fix(fleet): separate wake selection and remove capped device lookup

wakey-control-plane/src/api/control/fleet/mod.rs

@@ -6,7 +6,7 @@ use axum::http::StatusCode;
 use axum::response::IntoResponse;
 use futures_util::StreamExt;
 use tracing::{Instrument, info, warn};
-use wakey_agent::protocol::{AgentCommand, InventoryRequest, WakeRequest};
+use wakey_agent::protocol::{AgentCommand, InventoryRequest};
 
 use crate::api::ApiError;
 use crate::api::commands::{RelayCommandResponse, relay_agent_command};
@@ -14,6 +14,8 @@ use crate::runtime::AppState;
 
 mod build;
 mod types;
+mod wake;
+pub use wake::wake_fleet_device;
 
 #[cfg(test)]
 mod tests;
@@ -24,7 +26,7 @@ use build::{
 };
 use types::{
     FleetDevice, ListFleetDevicesQuery, RefreshFleetAgentResult, RefreshFleetDevicesRequest,
-    RefreshFleetDevicesResponse, WakeFleetDeviceRequest, WakeFleetDeviceResponse,
+    RefreshFleetDevicesResponse,
 };
 
 pub(crate) type FleetCache =
@@ -249,99 +251,6 @@ async fn refresh_fleet_agent_result(
     }
 }
 
-#[tracing::instrument(name = "fleet_wake", skip_all, fields(explicit_route = req.route_id.is_some(), outcome = tracing::field::Empty))]
-pub async fn wake_fleet_device(
-    State(state): State<AppState>,
-    Json(req): Json<WakeFleetDeviceRequest>,
-) -> Result<impl IntoResponse, ApiError> {
-    // Early database/relay errors retain an outcome even without a command result.
-    tracing::Span::current().record("outcome", "error");
-    let query = ListFleetDevicesQuery {
-        query: None,
-        presence: None,
-        known: None,
-        agent_id: None,
-        visibility: Some("all".into()),
-        limit: Some(1000),
-    };
-    let devices = load_fleet_devices(&state, &query).await.map_err(|err| {
-        warn!(error = %err, "failed loading fleet devices for wake");
-        fleet_error(err, "load_fleet_devices_failed")
-    })?;
-
-    let device = devices
-        .into_iter()
-        .find(|device| device.device_key == req.device_key)
-        .ok_or_else(|| {
-            tracing::Span::current().record("outcome", "device_not_found");
-            ApiError::new(
-                StatusCode::NOT_FOUND,
-                "fleet_device_not_found",
-                "fleet device not found",
-            )
-        })?;
-
-    let route = match req.route_id.as_deref() {
-        Some(route_id) => device
-            .route_candidates
-            .into_iter()
-            .find(|route| route.route_id == route_id),
-        None => device.recommended_route,
-    }
-    .ok_or_else(|| {
-        tracing::Span::current().record("outcome", "route_unavailable");
-        ApiError::new(
-            StatusCode::BAD_REQUEST,
-            "wake_route_unavailable",
-            "no wakeable connected MAC-backed route is available",
-        )
-    })?;
-
-    let Some(mac) = route.mac else {
-        tracing::Span::current().record("outcome", "route_without_mac");
-        return Err(ApiError::new(
-            StatusCode::BAD_REQUEST,
-            "wake_route_unavailable",
-            "selected route does not include a MAC address",
-        ));
-    };
-    if !route.connected {
-        tracing::Span::current().record("outcome", "route_disconnected");
-        return Err(ApiError::new(
-            StatusCode::BAD_REQUEST,
-            "wake_route_unavailable",
-            "selected route agent is not connected",
-        ));
-    }
-    if !route.wakeable {
-        tracing::Span::current().record("outcome", "route_not_wakeable");
-        return Err(ApiError::new(
-            StatusCode::BAD_REQUEST,
-            "wake_route_unavailable",
-            "selected route is not wakeable",
-        ));
-    }
-
-    info!(agent_id = %route.agent_id, has_ip = route.ip.is_some(), explicit_route = req.route_id.is_some(), "fleet wake route selected");
-    let command = relay_agent_command(
-        &state,
-        &route.agent_id,
-        AgentCommand::Wake(WakeRequest {
-            query: None,
-            mac: Some(mac),
-            ip: route.ip,
-        }),
-        req.timeout_ms,
-    )
-    .await?;
-    tracing::Span::current().record("outcome", command.status.as_str());
-
-    Ok((
-        StatusCode::OK,
-        Json(WakeFleetDeviceResponse { route, command }),
-    ))
-}
-
 #[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,

wakey-control-plane/src/api/control/fleet/wake.rs

@@ -0,0 +1,396 @@
+use axum::{Json, extract::State, http::StatusCode, response::IntoResponse};
+use macaddr::MacAddr;
+use tracing::{Instrument, info, warn};
+use wakey_agent::protocol::{AgentCommand, WakeRequest};
+
+use super::types::{FleetDevice, FleetWakeRoute, WakeFleetDeviceRequest, WakeFleetDeviceResponse};
+use super::{fleet_error, fleet_view};
+use crate::api::{ApiError, commands::relay_agent_command};
+use crate::runtime::AppState;
+
+#[tracing::instrument(name = "fleet_wake", skip_all, fields(explicit_route = req.route_id.is_some(), outcome = tracing::field::Empty))]
+pub async fn wake_fleet_device(
+    State(state): State<AppState>,
+    Json(req): Json<WakeFleetDeviceRequest>,
+) -> Result<impl IntoResponse, ApiError> {
+    // Early database/relay errors retain an outcome even without a command result.
+    tracing::Span::current().record("outcome", "error");
+    let selected = load_selected_route(&state, &req.device_key, req.route_id.as_deref()).await?;
+    info!(agent_id = %selected.route.agent_id, has_ip = selected.route.ip.is_some(),
+        explicit_route = req.route_id.is_some(), "fleet wake route selected");
+    let response = dispatch_wake(&state, selected, req.timeout_ms).await?;
+    tracing::Span::current().record("outcome", response.command.status.as_str());
+    Ok((StatusCode::OK, Json(response)))
+}
+
+/// List visibility and pagination are presentation policy, not wake eligibility.
+/// Borrow the complete cached view and clone only the selected route.
+async fn load_selected_route(
+    state: &AppState,
+    device_key: &str,
+    route_id: Option<&str>,
+) -> Result<SelectedRoute, ApiError> {
+    let span = tracing::info_span!("fleet_load", grouped = tracing::field::Empty);
+    let devices = fleet_view(state)
+        .instrument(span.clone())
+        .await
+        .map_err(|err| {
+            warn!(error = %err, "failed loading fleet devices for wake");
+            fleet_error(err, "load_fleet_devices_failed")
+        })?;
+    span.record("grouped", devices.len());
+    select_route(&devices, device_key, route_id).map_err(|error| {
+        tracing::Span::current().record("outcome", error.outcome());
+        error.into_api_error()
+    })
+}
+
+/// Selection does not rank routes: the fleet builder already chose the recommendation.
+/// An explicit route never falls back, including when it is disconnected or IP-only.
+fn select_route(
+    devices: &[FleetDevice],
+    device_key: &str,
+    route_id: Option<&str>,
+) -> Result<SelectedRoute, RouteError> {
+    let device = devices
+        .iter()
+        .find(|device| device.device_key == device_key)
+        .ok_or(RouteError::DeviceNotFound)?;
+    let route = match route_id {
+        Some(id) => device
+            .route_candidates
+            .iter()
+            .find(|route| route.route_id == id),
+        None => device.recommended_route.as_ref(),
+    }
+    .ok_or(RouteError::Unavailable)?;
+    let mac = route.mac.ok_or(RouteError::MissingMac)?;
+    if !route.connected {
+        return Err(RouteError::Disconnected);
+    }
+    if !route.wakeable {
+        return Err(RouteError::NotWakeable);
+    }
+    Ok(SelectedRoute {
+        route: route.clone(),
+        mac,
+    })
+}
+
+/// Only the selector constructs this; dispatch cannot accidentally use an IP-only route.
+#[derive(Debug)]
+struct SelectedRoute {
+    route: FleetWakeRoute,
+    mac: MacAddr,
+}
+
+impl SelectedRoute {
+    fn command(&self) -> AgentCommand {
+        AgentCommand::Wake(WakeRequest {
+            query: None,
+            mac: Some(self.mac),
+            ip: self.route.ip,
+        })
+    }
+}
+
+async fn dispatch_wake(
+    state: &AppState,
+    selected: SelectedRoute,
+    timeout_ms: Option<u64>,
+) -> Result<WakeFleetDeviceResponse, ApiError> {
+    // The relay checks the current agent connection again; cached connectivity
+    // is only selection evidence, not a guarantee that the agent is still there.
+    let command = relay_agent_command(
+        state,
+        &selected.route.agent_id,
+        selected.command(),
+        timeout_ms,
+    )
+    .await?;
+    Ok(WakeFleetDeviceResponse {
+        route: selected.route,
+        command,
+    })
+}
+
+#[derive(Debug, PartialEq, Eq)]
+enum RouteError {
+    DeviceNotFound,
+    Unavailable,
+    MissingMac,
+    Disconnected,
+    NotWakeable,
+}
+
+impl RouteError {
+    fn outcome(&self) -> &'static str {
+        match self {
+            Self::DeviceNotFound => "device_not_found",
+            Self::Unavailable => "route_unavailable",
+            Self::MissingMac => "route_without_mac",
+            Self::Disconnected => "route_disconnected",
+            Self::NotWakeable => "route_not_wakeable",
+        }
+    }
+
+    fn into_api_error(self) -> ApiError {
+        let message = match self {
+            Self::DeviceNotFound => {
+                return ApiError::new(
+                    StatusCode::NOT_FOUND,
+                    "fleet_device_not_found",
+                    "fleet device not found",
+                );
+            }
+            Self::Unavailable => "no wakeable connected MAC-backed route is available",
+            Self::MissingMac => "selected route does not include a MAC address",
+            Self::Disconnected => "selected route agent is not connected",
+            Self::NotWakeable => "selected route is not wakeable",
+        };
+        ApiError::new(StatusCode::BAD_REQUEST, "wake_route_unavailable", message)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::super::{load_fleet_devices, types::ListFleetDevicesQuery};
+    use super::*;
+    use wakey_core::Presence;
+
+    #[tokio::test]
+    async fn wake_lookup_can_select_a_device_beyond_the_list_limit() {
+        use std::collections::{BTreeSet, HashMap};
+        use std::sync::Arc;
+        use std::time::Duration;
+        use tokio::sync::{Mutex, RwLock};
+
+        let dir = std::env::temp_dir().join(format!("wakey-wake-limit-{}", uuid::Uuid::new_v4()));
+        let state = AppState {
+            store: Arc::new(
+                crate::state::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()),
+        };
+        // Seed the complete cached view, as the builder would. Exercise both
+        // production read paths without thousands of unrelated SQL inserts.
+        state
+            .fleet
+            .get_or_try_init(
+                || async { (state.store.fleet_revision(), BTreeSet::new()) },
+                || async {
+                    Ok((0..1001)
+                        .map(|i| {
+                            let mut row = device(vec![route("valid")], Some(route("valid")));
+                            row.device_key = format!("pc-{i}");
+                            row
+                        })
+                        .collect())
+                },
+            )
+            .await
+            .unwrap();
+
+        let listed = load_fleet_devices(
+            &state,
+            &ListFleetDevicesQuery {
+                visibility: Some("all".into()),
+                limit: Some(1000),
+                ..Default::default()
+            },
+        )
+        .await
+        .unwrap();
+        assert_eq!(listed.len(), 1000);
+        assert!(listed.iter().all(|row| row.device_key != "pc-1000"));
+        for route_id in [None, Some("valid")] {
+            let selected = load_selected_route(&state, "pc-1000", route_id)
+                .await
+                .unwrap();
+            assert_eq!(selected.route.route_id, "valid");
+        }
+        drop(state);
+        std::fs::remove_dir_all(dir).unwrap();
+    }
+
+    fn route(id: &str) -> FleetWakeRoute {
+        FleetWakeRoute {
+            route_id: id.into(),
+            agent_id: "agent-a".into(),
+            nickname: None,
+            connected: true,
+            mac: Some("aa:bb:cc:dd:ee:01".parse().unwrap()),
+            ip: Some("192.168.1.10".parse().unwrap()),
+            hostname: None,
+            interface: None,
+            source: "neighbor".into(),
+            presence: Presence::Offline,
+            last_seen_unix: 1,
+            wakeable: true,
+        }
+    }
+
+    fn device(candidates: Vec<FleetWakeRoute>, recommended: Option<FleetWakeRoute>) -> FleetDevice {
+        FleetDevice {
+            device_key: "pc".into(),
+            display_name: "pc".into(),
+            known_device: None,
+            pinned: false,
+            ips: vec![],
+            macs: vec![],
+            hostnames: vec![],
+            agents: vec![],
+            sources: vec![],
+            endpoints: vec![],
+            first_seen_unix: None,
+            last_seen_unix: None,
+            presence: Presence::Offline,
+            route_candidates: candidates,
+            recommended_route: recommended,
+        }
+    }
+
+    #[test]
+    fn selection_uses_recommendation_or_exact_override_and_preserves_address_pair() {
+        let first = route("first");
+        let mut recommended = route("recommended");
+        recommended.agent_id = "agent-b".into();
+        recommended.ip = Some("192.168.2.20".parse().unwrap());
+        let devices = [device(vec![first, recommended.clone()], Some(recommended))];
+        for (override_id, route_id, agent, ip) in [
+            (None, "recommended", "agent-b", "192.168.2.20"),
+            (Some("first"), "first", "agent-a", "192.168.1.10"),
+        ] {
+            let selected = select_route(&devices, "pc", override_id).unwrap();
+            assert_eq!(selected.route.route_id, route_id);
+            assert_eq!(selected.route.agent_id, agent);
+            let AgentCommand::Wake(command) = selected.command() else {
+                panic!("not a wake command")
+            };
+            assert!(command.query.is_none());
+            assert_eq!(command.mac, selected.route.mac);
+            assert_eq!(command.ip, Some(ip.parse().unwrap()));
+        }
+    }
+
+    #[test]
+    fn missing_device_and_missing_routes_keep_distinct_errors() {
+        let valid = route("valid");
+        let devices = [device(vec![valid.clone()], Some(valid))];
+        assert_eq!(
+            select_route(&devices, "missing", None).unwrap_err(),
+            RouteError::DeviceNotFound
+        );
+        assert_eq!(
+            select_route(&devices, "pc", Some("missing")).unwrap_err(),
+            RouteError::Unavailable
+        );
+        let devices = [device(vec![route("valid")], None)];
+        assert_eq!(
+            select_route(&devices, "pc", None).unwrap_err(),
+            RouteError::Unavailable
+        );
+    }
+
+    #[test]
+    fn invalid_explicit_routes_do_not_fall_back_and_validation_order_is_preserved() {
+        for (mac, connected, wakeable, expected) in [
+            (None, false, false, RouteError::MissingMac),
+            (
+                Some("aa:bb:cc:dd:ee:01"),
+                false,
+                false,
+                RouteError::Disconnected,
+            ),
+            (
+                Some("aa:bb:cc:dd:ee:01"),
+                true,
+                false,
+                RouteError::NotWakeable,
+            ),
+        ] {
+            let mut invalid = route("invalid");
+            invalid.mac = mac.map(|value| value.parse().unwrap());
+            invalid.connected = connected;
+            invalid.wakeable = wakeable;
+            let devices = [device(vec![invalid, route("valid")], Some(route("valid")))];
+            assert_eq!(
+                select_route(&devices, "pc", Some("invalid")).unwrap_err(),
+                expected
+            );
+        }
+    }
+
+    #[test]
+    fn mac_only_offline_route_is_still_wakeable() {
+        let mut mac_only = route("mac-only");
+        mac_only.ip = None;
+        let devices = [device(vec![mac_only.clone()], Some(mac_only))];
+        let selected = select_route(&devices, "pc", None).unwrap();
+        let AgentCommand::Wake(command) = selected.command() else {
+            panic!("not a wake command")
+        };
+        assert!(command.mac.is_some());
+        assert!(command.ip.is_none());
+    }
+
+    #[test]
+    fn route_errors_preserve_http_codes_messages_and_trace_outcomes() {
+        for (error, status, code, message, outcome) in [
+            (
+                RouteError::DeviceNotFound,
+                StatusCode::NOT_FOUND,
+                "fleet_device_not_found",
+                "fleet device not found",
+                "device_not_found",
+            ),
+            (
+                RouteError::Unavailable,
+                StatusCode::BAD_REQUEST,
+                "wake_route_unavailable",
+                "no wakeable connected MAC-backed route is available",
+                "route_unavailable",
+            ),
+            (
+                RouteError::MissingMac,
+                StatusCode::BAD_REQUEST,
+                "wake_route_unavailable",
+                "selected route does not include a MAC address",
+                "route_without_mac",
+            ),
+            (
+                RouteError::Disconnected,
+                StatusCode::BAD_REQUEST,
+                "wake_route_unavailable",
+                "selected route agent is not connected",
+                "route_disconnected",
+            ),
+            (
+                RouteError::NotWakeable,
+                StatusCode::BAD_REQUEST,
+                "wake_route_unavailable",
+                "selected route is not wakeable",
+                "route_not_wakeable",
+            ),
+        ] {
+            assert_eq!(error.outcome(), outcome);
+            let response = error.into_api_error();
+            assert_eq!(response.status, status);
+            assert_eq!(response.code, code);
+            assert_eq!(response.message, message);
+        }
+    }
+}