feat: mark agents unavailable from the dashboard for maintenance

kuiper-forge/migrations/shared/20261019000000_add_agent_unavailable_at.sql

@@ -0,0 +1,2 @@
+-- When an admin marked the agent unavailable (maintenance, image updates). NULL means it takes work.
+ALTER TABLE registered_agents ADD COLUMN unavailable_at TEXT;

kuiper-forge/src/admin/agent_log_routes.rs

@@ -82,6 +82,7 @@ pub(crate) async fn agent_logs(
         agent: AgentHeader {
             agent_id: agent.agent_id,
             revoked: agent.revoked,
+            unavailable_at: agent.unavailable_at,
         },
         tab: "logs",
         lines: lines

kuiper-forge/src/admin/api.rs

@@ -71,6 +71,8 @@ struct StatusResponse {
     registered_agents: usize,
     connected_agents: usize,
     revoked_agents: usize,
+    /// Marked unavailable by an admin: they get no new runners
+    unavailable_agents: usize,
     /// Sum of `max_vms` across connected agents
     total_capacity: usize,
     active_runners: usize,
@@ -99,6 +101,10 @@ async fn status(State(state): State<Arc<AdminState>>) -> Response {
         registered_agents: registered.len(),
         connected_agents: connected.len(),
         revoked_agents: registered.iter().filter(|a| a.revoked).count(),
+        unavailable_agents: registered
+            .iter()
+            .filter(|a| !a.revoked && a.unavailable_at.is_some())
+            .count(),
         total_capacity: connected.iter().map(|a| a.max_vms).sum(),
         active_runners: runners.len(),
         pending_jobs: pending_jobs.len(),
@@ -120,6 +126,10 @@ struct AgentResponse {
     hostname: String,
     agent_type: String,
     status: AgentStatus,
+    /// Marked unavailable by an admin (maintenance): gets no new runners, running ones finish
+    unavailable: bool,
+    /// When it was marked unavailable. Null while it takes work
+    unavailable_since: Option<DateTime<Utc>>,
     labels: Vec<String>,
     /// Only known while the agent is connected
     label_sets: Vec<Vec<String>>,
@@ -194,6 +204,8 @@ async fn agents(State(state): State<Arc<AdminState>>) -> Response {
                     .unwrap_or_default(),
                 last_seen_secs: live.map(|c| c.last_seen_secs),
                 status,
+                unavailable: a.unavailable_at.is_some(),
+                unavailable_since: a.unavailable_at,
                 agent_id: a.agent_id,
                 hostname: a.hostname,
                 agent_type: a.agent_type,

kuiper-forge/src/admin/middleware.rs

@@ -6,6 +6,7 @@ use crate::agent_failures::AgentFailureStore;
 use crate::agent_logs::AgentLogStore;
 use crate::agent_registry::AgentRegistry;
 use crate::auth::AuthManager;
+use crate::fleet::FleetNotifier;
 use crate::pending_jobs::PendingJobStore;
 use crate::runner_state::RunnerStateStore;
 use crate::tls::ServerTrust;
@@ -38,4 +39,6 @@ pub struct AdminState {
     pub server_trust: ServerTrust,
     /// Coordinator URL for registration bundles
     pub coordinator_url: String,
+    /// Wakes the fleet manager, e.g. when an agent takes work again
+    pub fleet_notifier: Option<FleetNotifier>,
 }

kuiper-forge/src/admin/routes.rs

@@ -25,7 +25,7 @@ use chrono::{Duration, Utc};
 use serde::Deserialize;
 use std::collections::HashMap;
 use std::sync::Arc;
-use tracing::error;
+use tracing::{error, info, warn};
 
 /// Build the admin router.
 pub fn admin_router(state: Arc<AdminState>) -> Router {
@@ -38,6 +38,7 @@ pub fn admin_router(state: Arc<AdminState>) -> Router {
         .route("/tokens/{token}/delete", post(token_delete))
         .route("/agents/{agent_id}", get(agent_detail))
         .route("/agents/{agent_id}/revoke", post(agent_revoke))
+        .route("/agents/{agent_id}/availability", post(agent_availability))
         .route("/agents/{agent_id}/logs", get(agent_log_routes::agent_logs))
         .route("/users", get(user_routes::users_page))
         .route("/users/create", post(user_routes::user_create))
@@ -200,15 +201,22 @@ async fn render_dashboard(
     let now = Utc::now();
     let mut pending_jobs = Vec::with_capacity(pending.len());
     for (job_id, job) in pending {
+        let matching = state
+            .agent_registry
+            .find_agents_by_labels(&job.agent_labels)
+            .await;
+        let mut unavailable_agents = 0;
+        for agent_id in &matching {
+            if state.agent_registry.is_unavailable(agent_id).await {
+                unavailable_agents += 1;
+            }
+        }
         pending_jobs.push(PendingJobSummary {
             assigned_agent: job_agents.get(&job_id).cloned(),
             job_id,
             waiting: format_age(now - job.created_at),
-            matching_agents: state
-                .agent_registry
-                .find_agents_by_labels(&job.agent_labels)
-                .await
-                .len(),
+            matching_agents: matching.len(),
+            unavailable_agents,
             free_capacity: state
                 .agent_registry
                 .available_capacity(&job.agent_labels)
@@ -252,6 +260,7 @@ async fn render_dashboard(
                 active_vms,
                 created_at: a.created_at,
                 revoked: a.revoked,
+                unavailable_at: a.unavailable_at,
                 version: a.agent_version,
             }
         })
@@ -437,6 +446,7 @@ async fn agent_detail(
         active_vms: runners.len(),
         created_at: agent.created_at,
         revoked: agent.revoked,
+        unavailable_at: agent.unavailable_at,
         version: agent.agent_version,
     };
 
@@ -496,6 +506,83 @@ async fn agent_revoke(
     }
 }
 
+/// Availability toggle form data.
+#[derive(Deserialize)]
+pub struct AvailabilityForm {
+    /// `true` to take work again, `false` for maintenance
+    available: bool,
+    /// Page to go back to: `dashboard`, else the agent's page
+    #[serde(default)]
+    from: String,
+}
+
+/// Mark an agent unavailable (no new runners, running ones finish) or available again.
+async fn agent_availability(
+    State(state): State<Arc<AdminState>>,
+    jar: CookieJar,
+    Path(agent_id): Path<String>,
+    Form(form): Form<AvailabilityForm>,
+) -> Response {
+    let session = match check_auth(&state, &jar).await {
+        Some(s) => s,
+        None => return Redirect::to("/admin/login").into_response(),
+    };
+
+    // DB first: if that fails, scheduling keeps doing what the UI shows
+    match state
+        .auth_manager
+        .set_agent_unavailable(&agent_id, !form.available)
+        .await
+    {
+        Ok(true) => {}
+        Ok(false) => return (StatusCode::NOT_FOUND, "Agent not found").into_response(),
+        Err(e) => {
+            error!(
+                agent_id = %agent_id,
+                available = form.available,
+                "Failed to update agent availability: {:#}", e
+            );
+            return (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Failed to update agent availability",
+            )
+                .into_response();
+        }
+    }
+    state
+        .agent_registry
+        .set_unavailable(&agent_id, !form.available)
+        .await;
+
+    let active_runners = state
+        .runner_state
+        .get_all_runners()
+        .await
+        .map(|r| r.iter().filter(|(_, r)| r.agent_id == agent_id).count())
+        .unwrap_or_else(|e| {
+            warn!(agent_id = %agent_id, "Failed to count runners: {}", e);
+            0
+        });
+    info!(
+        agent_id = %agent_id,
+        available = form.available,
+        active_runners = active_runners,
+        by = %session.username,
+        "Agent availability set from admin UI"
+    );
+
+    // right away instead of on the next tick: queued jobs and pool top-ups, or draining idle pool runners
+    if let Some(notifier) = &state.fleet_notifier {
+        notifier.notify().await;
+    }
+
+    if form.from == "dashboard" {
+        Redirect::to("/admin/dashboard").into_response()
+    } else {
+        Redirect::to(&format!("/admin/agents/{agent_id}")).into_response()
+    }
+}
+
 /// Short human age like `45s`, `4m 12s`, `3h 5m` or `2d 4h`.
 fn format_age(age: Duration) -> String {
     let secs = age.num_seconds().max(0);

kuiper-forge/src/admin/templates.rs

@@ -52,6 +52,8 @@ pub struct PendingJobSummary {
     pub failed_agents: usize,
     /// Connected agents whose labels match, regardless of free capacity
     pub matching_agents: usize,
+    /// Of `matching_agents`, the ones marked unavailable
+    pub unavailable_agents: usize,
     pub free_capacity: usize,
 }
 
@@ -71,6 +73,8 @@ pub struct AgentSummary {
     pub active_vms: usize,
     pub created_at: DateTime<Utc>,
     pub revoked: bool,
+    /// When an admin marked it unavailable (no new runners). None while it takes work
+    pub unavailable_at: Option<DateTime<Utc>>,
     /// None for agents that predate version reporting
     pub version: Option<String>,
 }
@@ -132,6 +136,7 @@ pub struct AgentDetailTemplate {
 pub struct AgentHeader {
     pub agent_id: String,
     pub revoked: bool,
+    pub unavailable_at: Option<DateTime<Utc>>,
 }
 
 /// Agent logs tab
@@ -242,6 +247,7 @@ mod tests {
             limits: vec![],
             explicit_pools: false,
             last_seen_secs: 0,
+            unavailable: false,
         }
     }