what

.github/copilot-instructions.md

@@ -1 +1,3 @@
 We work towards clean, modular, maintainable design, clear documentation... im not the best at this thing, so you help me.
+
+DRY: spawn subagents. 

wakey-control-plane/src/api.rs

@@ -0,0 +1,185 @@
+use axum::Json;
+use axum::extract::{Path as AxumPath, State};
+use axum::http::StatusCode;
+use axum::response::IntoResponse;
+use serde::{Deserialize, Serialize};
+use uuid::Uuid;
+use wakey_agent::protocol::{AgentCommand, ErrorPayload, RequestId, ServerMessage};
+
+use crate::runtime::{AgentReply, AppState};
+
+#[derive(Debug, Deserialize)]
+pub struct EnrollRequest {
+    pub enroll_token: String,
+}
+
+#[derive(Debug, Serialize)]
+pub struct EnrollResponse {
+    pub agent_id: String,
+    pub agent_token: String,
+    pub server_url: String,
+}
+
+#[derive(Debug, Serialize)]
+pub struct AgentStatus {
+    pub agent_id: String,
+    pub connected: bool,
+}
+
+#[derive(Debug, Deserialize)]
+pub struct RelayCommandRequest {
+    pub command: AgentCommand,
+    pub timeout_ms: Option<u64>,
+}
+
+#[derive(Debug, Serialize)]
+pub struct RelayCommandResponse {
+    pub request_id: String,
+    pub status: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub result: Option<serde_json::Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub error: Option<ErrorPayload>,
+}
+
+pub async fn healthz() -> &'static str {
+    "ok"
+}
+
+pub async fn enroll(
+    State(state): State<AppState>,
+    Json(req): Json<EnrollRequest>,
+) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
+    match state.store.enroll(&req.enroll_token).await {
+        Ok(issued) => Ok((
+            StatusCode::OK,
+            Json(EnrollResponse {
+                agent_id: issued.agent_id,
+                agent_token: issued.agent_token,
+                server_url: state.public_url,
+            }),
+        )),
+        Err(err) => Err(json_error(
+            StatusCode::UNAUTHORIZED,
+            "enrollment_rejected",
+            &err.to_string(),
+        )),
+    }
+}
+
+pub async fn list_agents(
+    State(state): State<AppState>,
+) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
+    let enrolled = state.store.list_agents().await;
+    let sessions = state.sessions.read().await;
+
+    let agents = enrolled
+        .into_iter()
+        .map(|agent_id| AgentStatus {
+            connected: sessions.contains_key(&agent_id),
+            agent_id,
+        })
+        .collect::<Vec<_>>();
+
+    Ok((StatusCode::OK, Json(agents)))
+}
+
+pub async fn run_command(
+    State(state): State<AppState>,
+    AxumPath(agent_id): AxumPath<String>,
+    Json(req): Json<RelayCommandRequest>,
+) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
+    let request_id_string = format!("req-{}", Uuid::new_v4());
+    let request_id = RequestId::try_from(request_id_string.clone()).map_err(|err| {
+        json_error(
+            StatusCode::INTERNAL_SERVER_ERROR,
+            "invalid_request_id",
+            &err,
+        )
+    })?;
+
+    let tx = {
+        let sessions = state.sessions.read().await;
+        sessions.get(&agent_id).cloned()
+    }
+    .ok_or_else(|| {
+        json_error(
+            StatusCode::NOT_FOUND,
+            "agent_not_connected",
+            "agent is not connected",
+        )
+    })?;
+
+    let (pending_tx, pending_rx) = tokio::sync::oneshot::channel();
+    state
+        .pending
+        .lock()
+        .await
+        .insert(request_id_string.clone(), pending_tx);
+
+    if let Err(err) = tx.send(ServerMessage::Command {
+        request_id,
+        command: req.command,
+    }) {
+        state.pending.lock().await.remove(&request_id_string);
+        return Err(json_error(
+            StatusCode::BAD_GATEWAY,
+            "agent_send_failed",
+            &format!("failed to send command to agent: {err}"),
+        ));
+    }
+
+    let timeout = std::time::Duration::from_millis(
+        req.timeout_ms
+            .unwrap_or(state.command_timeout.as_millis() as u64)
+            .max(1),
+    );
+    let outcome = tokio::time::timeout(timeout, pending_rx).await;
+    let response = match outcome {
+        Ok(Ok(AgentReply::Result(result))) => RelayCommandResponse {
+            request_id: request_id_string,
+            status: "ok".into(),
+            result: Some(result),
+            error: None,
+        },
+        Ok(Ok(AgentReply::Error(error))) => RelayCommandResponse {
+            request_id: request_id_string,
+            status: "error".into(),
+            result: None,
+            error: Some(error),
+        },
+        Ok(Err(_)) => {
+            return Err(json_error(
+                StatusCode::BAD_GATEWAY,
+                "agent_response_dropped",
+                "agent response channel dropped",
+            ));
+        }
+        Err(_) => {
+            state.pending.lock().await.remove(&request_id_string);
+            return Err(json_error(
+                StatusCode::GATEWAY_TIMEOUT,
+                "agent_timeout",
+                "agent did not answer before timeout",
+            ));
+        }
+    };
+
+    Ok((StatusCode::OK, Json(response)))
+}
+
+pub fn json_error(
+    status: StatusCode,
+    code: &str,
+    message: &str,
+) -> (StatusCode, Json<serde_json::Value>) {
+    (
+        status,
+        Json(serde_json::json!({
+            "error": {
+                "code": code,
+                "message": message,
+            }
+        })),
+    )
+}

wakey-control-plane/src/cli.rs

@@ -0,0 +1,67 @@
+use std::net::SocketAddr;
+use std::path::PathBuf;
+
+use clap::{ArgAction, Args, Parser, Subcommand};
+
+pub const DEFAULT_STATE_FILE: &str = "/var/lib/wakey-control-plane/state.json";
+pub const DEFAULT_PID_FILE: &str = "/var/run/wakey-control-plane.pid";
+
+#[derive(Parser)]
+#[command(name = "wakey-control-plane")]
+#[command(version, about = "Control plane server for wakey-agent fleets")]
+pub struct Cli {
+    #[arg(short = 'v', long = "verbose", action = ArgAction::Count, global = true)]
+    pub verbose: u8,
+
+    #[command(subcommand)]
+    pub command: Command,
+}
+
+#[derive(Subcommand)]
+pub enum Command {
+    /// Run the control-plane daemon.
+    Serve(ServeArgs),
+    /// Create a new enroll token for provisioning a router.
+    IssueEnrollToken(IssueEnrollTokenArgs),
+    /// Send SIGHUP to an already-running daemon.
+    Reload(ReloadArgs),
+}
+
+#[derive(Args, Clone)]
+pub struct ServeArgs {
+    #[arg(long, default_value = "0.0.0.0:8080")]
+    pub bind: SocketAddr,
+
+    #[arg(long, default_value = "http://127.0.0.1:8080")]
+    pub public_url: String,
+
+    #[arg(long, default_value = DEFAULT_STATE_FILE)]
+    pub state_file: PathBuf,
+
+    #[arg(long = "enroll-token")]
+    pub enroll_tokens: Vec<String>,
+
+    #[arg(long, default_value_t = 30_000)]
+    pub command_timeout_ms: u64,
+
+    #[arg(long, default_value = DEFAULT_PID_FILE)]
+    pub pid_file: PathBuf,
+}
+
+#[derive(Args)]
+pub struct IssueEnrollTokenArgs {
+    #[arg(long, default_value = DEFAULT_STATE_FILE)]
+    pub state_file: PathBuf,
+
+    #[arg(long = "enroll-token")]
+    pub enroll_tokens: Vec<String>,
+
+    #[arg(long)]
+    pub public_url: Option<String>,
+}
+
+#[derive(Args)]
+pub struct ReloadArgs {
+    #[arg(long, default_value = DEFAULT_PID_FILE)]
+    pub pid_file: PathBuf,
+}

wakey-control-plane/src/main.rs

@@ -1,416 +1,22 @@
+mod api;
+mod cli;
+mod runtime;
 mod state;
 mod tracing;
+mod ws;
 
-use std::collections::HashMap;
-use std::net::SocketAddr;
-use std::sync::Arc;
-use std::time::Duration;
-
-use ::tracing::{debug, info, warn};
-use anyhow::{Context, Result};
-use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
-use axum::extract::{Path, State};
-use axum::http::StatusCode;
-use axum::response::IntoResponse;
-use axum::routing::{get, post};
-use axum::{Json, Router};
-use clap::{ArgAction, Parser};
-use futures_util::{SinkExt, StreamExt};
-use serde::{Deserialize, Serialize};
-use tokio::net::TcpListener;
-use tokio::sync::{Mutex, RwLock, mpsc, oneshot};
-use uuid::Uuid;
-use wakey_agent::protocol::{AgentCommand, ErrorPayload, RequestId, ServerMessage};
-
-#[derive(Parser)]
-#[command(name = "wakey-control-plane")]
-#[command(version, about = "Control plane server for wakey-agent fleets")]
-struct Cli {
-    #[arg(short = 'v', long = "verbose", action = ArgAction::Count, global = true)]
-    verbose: u8,
-
-    #[arg(long, default_value = "0.0.0.0:8080")]
-    bind: SocketAddr,
-
-    #[arg(long, default_value = "http://127.0.0.1:8080")]
-    public_url: String,
-
-    #[arg(long, default_value = "/var/lib/wakey-control-plane/state.json")]
-    state_file: std::path::PathBuf,
-
-    #[arg(long = "enroll-token")]
-    enroll_tokens: Vec<String>,
-
-    #[arg(long, default_value_t = 30_000)]
-    command_timeout_ms: u64,
-}
-
-#[derive(Clone)]
-struct AppState {
-    store: Arc<state::Store>,
-    sessions: Arc<RwLock<HashMap<String, mpsc::UnboundedSender<ServerMessage>>>>,
-    pending: Arc<Mutex<HashMap<String, oneshot::Sender<AgentReply>>>>,
-    public_url: String,
-    command_timeout: Duration,
-}
-
-enum AgentReply {
-    Result(serde_json::Value),
-    Error(ErrorPayload),
-}
-
-#[derive(Debug, Deserialize)]
-#[serde(tag = "type", rename_all = "snake_case")]
-enum IncomingClientMessage {
-    Hello {
-        agent_id: String,
-    },
-    Auth {
-        agent_id: String,
-        agent_token: String,
-    },
-    Heartbeat {
-        agent_id: String,
-    },
-    Result {
-        request_id: RequestId,
-        result: serde_json::Value,
-    },
-    Error {
-        request_id: RequestId,
-        error: ErrorPayload,
-    },
-}
-
-#[derive(Debug, Deserialize)]
-struct EnrollRequest {
-    enroll_token: String,
-}
-
-#[derive(Debug, Serialize)]
-struct EnrollResponse {
-    agent_id: String,
-    agent_token: String,
-    server_url: String,
-}
-
-#[derive(Debug, Serialize)]
-struct AgentStatus {
-    agent_id: String,
-    connected: bool,
-}
-
-#[derive(Debug, Deserialize)]
-struct RelayCommandRequest {
-    command: AgentCommand,
-    timeout_ms: Option<u64>,
-}
-
-#[derive(Debug, Serialize)]
-struct RelayCommandResponse {
-    request_id: String,
-    status: String,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    result: Option<serde_json::Value>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    error: Option<ErrorPayload>,
-}
+use anyhow::Result;
+use clap::Parser;
+use cli::{Cli, Command};
 
 #[tokio::main]
 async fn main() -> Result<()> {
     let cli = Cli::parse();
     tracing::init(cli.verbose);
 
-    let store = state::Store::load_or_init(&cli.state_file, cli.enroll_tokens)
-        .await
-        .with_context(|| format!("failed to initialize store {}", cli.state_file.display()))?;
-
-    let app_state = AppState {
-        store: Arc::new(store),
-        sessions: Arc::new(RwLock::new(HashMap::new())),
-        pending: Arc::new(Mutex::new(HashMap::new())),
-        public_url: cli.public_url.trim_end_matches('/').to_string(),
-        command_timeout: Duration::from_millis(cli.command_timeout_ms.max(1)),
-    };
-
-    let app = Router::new()
-        .route("/healthz", get(healthz))
-        .route("/api/v1/agents/enroll", post(enroll))
-        .route("/api/v1/agent/ws", get(agent_ws))
-        .route("/api/v1/control/agents", get(list_agents))
-        .route(
-            "/api/v1/control/agents/{agent_id}/command",
-            post(run_command),
-        )
-        .with_state(app_state);
-
-    info!(bind = %cli.bind, "starting control-plane server");
-    let listener = TcpListener::bind(cli.bind).await?;
-    axum::serve(listener, app).await?;
-    Ok(())
-}
-
-async fn healthz() -> &'static str {
-    "ok"
-}
-
-async fn enroll(
-    State(state): State<AppState>,
-    Json(req): Json<EnrollRequest>,
-) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
-    match state.store.enroll(&req.enroll_token).await {
-        Ok(issued) => Ok((
-            StatusCode::OK,
-            Json(EnrollResponse {
-                agent_id: issued.agent_id,
-                agent_token: issued.agent_token,
-                server_url: state.public_url,
-            }),
-        )),
-        Err(err) => Err(json_error(
-            StatusCode::UNAUTHORIZED,
-            "enrollment_rejected",
-            &err.to_string(),
-        )),
+    match cli.command {
+        Command::Serve(args) => runtime::serve(args).await,
+        Command::IssueEnrollToken(args) => runtime::issue_enroll_token(args).await,
+        Command::Reload(args) => runtime::reload_daemon(&args.pid_file),
     }
 }
-
-async fn list_agents(
-    State(state): State<AppState>,
-) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
-    let enrolled = state.store.list_agents().await;
-    let sessions = state.sessions.read().await;
-
-    let agents = enrolled
-        .into_iter()
-        .map(|agent_id| AgentStatus {
-            connected: sessions.contains_key(&agent_id),
-            agent_id,
-        })
-        .collect::<Vec<_>>();
-
-    Ok((StatusCode::OK, Json(agents)))
-}
-
-async fn run_command(
-    State(state): State<AppState>,
-    Path(agent_id): Path<String>,
-    Json(req): Json<RelayCommandRequest>,
-) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
-    let request_id_string = format!("req-{}", Uuid::new_v4());
-    let request_id = RequestId::try_from(request_id_string.clone()).map_err(|err| {
-        json_error(
-            StatusCode::INTERNAL_SERVER_ERROR,
-            "invalid_request_id",
-            &err,
-        )
-    })?;
-
-    let tx = {
-        let sessions = state.sessions.read().await;
-        sessions.get(&agent_id).cloned()
-    }
-    .ok_or_else(|| {
-        json_error(
-            StatusCode::NOT_FOUND,
-            "agent_not_connected",
-            "agent is not connected",
-        )
-    })?;
-
-    let (pending_tx, pending_rx) = oneshot::channel();
-    state
-        .pending
-        .lock()
-        .await
-        .insert(request_id_string.clone(), pending_tx);
-
-    if let Err(err) = tx.send(ServerMessage::Command {
-        request_id,
-        command: req.command,
-    }) {
-        state.pending.lock().await.remove(&request_id_string);
-        return Err(json_error(
-            StatusCode::BAD_GATEWAY,
-            "agent_send_failed",
-            &format!("failed to send command to agent: {err}"),
-        ));
-    }
-
-    let timeout = Duration::from_millis(
-        req.timeout_ms
-            .unwrap_or(state.command_timeout.as_millis() as u64)
-            .max(1),
-    );
-    let outcome = tokio::time::timeout(timeout, pending_rx).await;
-    let response = match outcome {
-        Ok(Ok(AgentReply::Result(result))) => RelayCommandResponse {
-            request_id: request_id_string,
-            status: "ok".into(),
-            result: Some(result),
-            error: None,
-        },
-        Ok(Ok(AgentReply::Error(error))) => RelayCommandResponse {
-            request_id: request_id_string,
-            status: "error".into(),
-            result: None,
-            error: Some(error),
-        },
-        Ok(Err(_)) => {
-            return Err(json_error(
-                StatusCode::BAD_GATEWAY,
-                "agent_response_dropped",
-                "agent response channel dropped",
-            ));
-        }
-        Err(_) => {
-            state.pending.lock().await.remove(&request_id_string);
-            return Err(json_error(
-                StatusCode::GATEWAY_TIMEOUT,
-                "agent_timeout",
-                "agent did not answer before timeout",
-            ));
-        }
-    };
-
-    Ok((StatusCode::OK, Json(response)))
-}
-
-async fn agent_ws(ws: WebSocketUpgrade, State(state): State<AppState>) -> impl IntoResponse {
-    ws.on_upgrade(move |socket| handle_agent_socket(state, socket))
-}
-
-async fn handle_agent_socket(state: AppState, socket: WebSocket) {
-    let (mut write, mut read) = socket.split();
-    let (tx, mut rx) = mpsc::unbounded_channel::<ServerMessage>();
-
-    let writer = tokio::spawn(async move {
-        while let Some(msg) = rx.recv().await {
-            let encoded = match serde_json::to_string(&msg) {
-                Ok(s) => s,
-                Err(err) => {
-                    warn!(error = %err, "failed to encode server websocket message");
-                    continue;
-                }
-            };
-            if let Err(err) = write.send(Message::Text(encoded.into())).await {
-                warn!(error = %err, "failed to send websocket message");
-                break;
-            }
-        }
-    });
-
-    let mut authed_agent_id: Option<String> = None;
-
-    loop {
-        let frame = read.next().await;
-        let msg = match frame {
-            Some(Ok(msg)) => msg,
-            Some(Err(err)) => {
-                warn!(error = %err, "agent websocket receive error");
-                break;
-            }
-            None => break,
-        };
-
-        match msg {
-            Message::Text(text) => {
-                if let Err(err) = process_agent_text(&state, &tx, &mut authed_agent_id, &text).await
-                {
-                    warn!(error = %err, "closing agent websocket due to protocol/auth error");
-                    break;
-                }
-            }
-            Message::Ping(_) => {}
-            Message::Pong(_) => {}
-            Message::Close(_) => break,
-            Message::Binary(_) => {
-                debug!("ignoring unexpected binary websocket frame");
-            }
-        }
-    }
-
-    if let Some(agent_id) = authed_agent_id {
-        info!(agent_id = %agent_id, "agent disconnected");
-        state.sessions.write().await.remove(&agent_id);
-    }
-
-    writer.abort();
-}
-
-async fn process_agent_text(
-    state: &AppState,
-    tx: &mpsc::UnboundedSender<ServerMessage>,
-    authed_agent_id: &mut Option<String>,
-    text: &str,
-) -> Result<()> {
-    let message: IncomingClientMessage =
-        serde_json::from_str(text).context("invalid client websocket payload")?;
-
-    match message {
-        IncomingClientMessage::Hello { agent_id } => {
-            debug!(agent_id = %agent_id, "agent hello received");
-        }
-        IncomingClientMessage::Auth {
-            agent_id,
-            agent_token,
-        } => {
-            if !state
-                .store
-                .verify_agent_token(&agent_id, &agent_token)
-                .await
-            {
-                anyhow::bail!("agent auth rejected");
-            }
-            state
-                .sessions
-                .write()
-                .await
-                .insert(agent_id.clone(), tx.clone());
-            *authed_agent_id = Some(agent_id.clone());
-            info!(agent_id = %agent_id, "agent authenticated");
-        }
-        IncomingClientMessage::Heartbeat { agent_id } => {
-            if authed_agent_id.as_deref() != Some(agent_id.as_str()) {
-                anyhow::bail!("heartbeat for unauthenticated or mismatched agent");
-            }
-            debug!(agent_id = %agent_id, "heartbeat received");
-        }
-        IncomingClientMessage::Result { request_id, result } => {
-            if authed_agent_id.is_none() {
-                anyhow::bail!("result before auth");
-            }
-            let key = request_id.as_str().to_string();
-            if let Some(waiter) = state.pending.lock().await.remove(&key) {
-                let _ = waiter.send(AgentReply::Result(result));
-            }
-        }
-        IncomingClientMessage::Error { request_id, error } => {
-            if authed_agent_id.is_none() {
-                anyhow::bail!("error before auth");
-            }
-            let key = request_id.as_str().to_string();
-            if let Some(waiter) = state.pending.lock().await.remove(&key) {
-                let _ = waiter.send(AgentReply::Error(error));
-            }
-        }
-    }
-
-    Ok(())
-}
-
-fn json_error(
-    status: StatusCode,
-    code: &str,
-    message: &str,
-) -> (StatusCode, Json<serde_json::Value>) {
-    (
-        status,
-        Json(serde_json::json!({
-            "error": {
-                "code": code,
-                "message": message,
-            }
-        })),
-    )
-}

wakey-control-plane/src/runtime.rs

@@ -0,0 +1,165 @@
+use std::collections::HashMap;
+use std::path::Path;
+use std::sync::Arc;
+use std::time::Duration;
+
+use anyhow::{Context, Result};
+use axum::Router;
+use axum::routing::{get, post};
+use tokio::net::TcpListener;
+use tokio::sync::{Mutex, RwLock, mpsc, oneshot};
+use tracing::{info, warn};
+use wakey_agent::protocol::{ErrorPayload, ServerMessage};
+
+use crate::api;
+use crate::cli::{IssueEnrollTokenArgs, ServeArgs};
+use crate::state;
+use crate::ws;
+
+#[derive(Clone)]
+pub struct AppState {
+    pub store: Arc<state::Store>,
+    pub sessions: Arc<RwLock<HashMap<String, mpsc::UnboundedSender<ServerMessage>>>>,
+    pub pending: Arc<Mutex<HashMap<String, oneshot::Sender<AgentReply>>>>,
+    pub public_url: String,
+    pub command_timeout: Duration,
+}
+
+pub enum AgentReply {
+    Result(serde_json::Value),
+    Error(ErrorPayload),
+}
+
+pub async fn serve(args: ServeArgs) -> Result<()> {
+    write_pid_file(&args.pid_file)?;
+
+    let store = state::Store::load_or_init(&args.state_file, args.enroll_tokens)
+        .await
+        .with_context(|| format!("failed to initialize store {}", args.state_file.display()))?;
+
+    let app_state = AppState {
+        store: Arc::new(store),
+        sessions: Arc::new(RwLock::new(HashMap::new())),
+        pending: Arc::new(Mutex::new(HashMap::new())),
+        public_url: args.public_url.trim_end_matches('/').to_string(),
+        command_timeout: Duration::from_millis(args.command_timeout_ms.max(1)),
+    };
+
+    let app = Router::new()
+        .route("/healthz", get(api::healthz))
+        .route("/api/v1/agents/enroll", post(api::enroll))
+        .route("/api/v1/agent/ws", get(ws::agent_ws))
+        .route("/api/v1/control/agents", get(api::list_agents))
+        .route(
+            "/api/v1/control/agents/{agent_id}/command",
+            post(api::run_command),
+        )
+        .with_state(app_state.clone());
+
+    info!(bind = %args.bind, pid_file = %args.pid_file.display(), "starting control-plane server");
+    let listener = TcpListener::bind(args.bind).await?;
+    let mut server = tokio::spawn(async move {
+        axum::serve(listener, app)
+            .await
+            .context("control-plane server exited unexpectedly")
+    });
+
+    #[cfg(unix)]
+    {
+        use tokio::signal::unix::{SignalKind, signal};
+        let mut hup = signal(SignalKind::hangup()).context("failed to install SIGHUP handler")?;
+
+        loop {
+            tokio::select! {
+                _ = tokio::signal::ctrl_c() => {
+                    info!("ctrl-c received; shutting down control-plane");
+                    server.abort();
+                    break;
+                }
+                _ = hup.recv() => {
+                    match app_state.store.reload_from_disk().await {
+                        Ok(()) => info!("reloaded state from disk"),
+                        Err(err) => warn!(error = %err, "failed to reload state from disk"),
+                    }
+                }
+                join = &mut server => {
+                    let _ = remove_pid_file(&args.pid_file);
+                    return join.context("control-plane join failed")?;
+                }
+            }
+        }
+    }
+
+    #[cfg(not(unix))]
+    {
+        tokio::signal::ctrl_c()
+            .await
+            .context("failed waiting for ctrl-c")?;
+        server.abort();
+    }
+
+    let _ = remove_pid_file(&args.pid_file);
+    Ok(())
+}
+
+pub async fn issue_enroll_token(args: IssueEnrollTokenArgs) -> Result<()> {
+    let store = state::Store::load_or_init(&args.state_file, args.enroll_tokens)
+        .await
+        .with_context(|| format!("failed to initialize store {}", args.state_file.display()))?;
+    let token = store.issue_enroll_token().await?;
+    println!("enroll_token={token}");
+    if let Some(url) = args.public_url {
+        let base = url.trim_end_matches('/');
+        println!("agent_command=wakey-agent enroll --server-url {base} --enroll-token {token}");
+    }
+    Ok(())
+}
+
+pub fn reload_daemon(pid_file: &Path) -> Result<()> {
+    let pid = read_pid(pid_file)?;
+    send_hup(pid)
+}
+
+fn write_pid_file(path: &Path) -> Result<()> {
+    if let Some(parent) = path.parent() {
+        std::fs::create_dir_all(parent)
+            .with_context(|| format!("failed to create pid dir {}", parent.display()))?;
+    }
+    std::fs::write(path, format!("{}\n", std::process::id()))
+        .with_context(|| format!("failed to write pid file {}", path.display()))
+}
+
+fn remove_pid_file(path: &Path) -> Result<()> {
+    match std::fs::remove_file(path) {
+        Ok(()) => Ok(()),
+        Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()),
+        Err(err) => {
+            Err(err).with_context(|| format!("failed to remove pid file {}", path.display()))
+        }
+    }
+}
+
+fn read_pid(path: &Path) -> Result<i32> {
+    let raw = std::fs::read_to_string(path)
+        .with_context(|| format!("failed to read pid file {}", path.display()))?;
+    let pid = raw
+        .trim()
+        .parse::<i32>()
+        .with_context(|| format!("invalid pid in {}", path.display()))?;
+    if pid <= 0 {
+        anyhow::bail!("invalid non-positive pid {pid}");
+    }
+    Ok(pid)
+}
+
+fn send_hup(pid: i32) -> Result<()> {
+    let status = std::process::Command::new("kill")
+        .arg("-HUP")
+        .arg(pid.to_string())
+        .status()
+        .context("failed to invoke kill -HUP")?;
+    if !status.success() {
+        anyhow::bail!("kill -HUP failed for pid {pid}");
+    }
+    Ok(())
+}

wakey-control-plane/src/state.rs

@@ -70,6 +70,30 @@ impl Store {
         })
     }
 
+    pub async fn issue_enroll_token(&self) -> Result<String> {
+        let token = format!("enr-{}", Uuid::new_v4());
+        let mut state = self.state.write().await;
+        state.enroll_tokens.insert(token.clone());
+        drop(state);
+        self.save().await?;
+        Ok(token)
+    }
+
+    pub async fn reload_from_disk(&self) -> Result<()> {
+        if !self.path.exists() {
+            return Ok(());
+        }
+
+        let raw = tokio::fs::read_to_string(&self.path)
+            .await
+            .with_context(|| format!("failed to read store {}", self.path.display()))?;
+        let decoded = serde_json::from_str::<PersistedState>(&raw)
+            .with_context(|| format!("failed to decode store {}", self.path.display()))?;
+
+        *self.state.write().await = decoded;
+        Ok(())
+    }
+
     pub async fn verify_agent_token(&self, agent_id: &str, token: &str) -> bool {
         self.state
             .read()

wakey-control-plane/src/ws.rs

@@ -0,0 +1,157 @@
+use anyhow::{Context, Result};
+use axum::extract::State;
+use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
+use axum::response::IntoResponse;
+use futures_util::{SinkExt, StreamExt};
+use serde::Deserialize;
+use tokio::sync::mpsc;
+use tracing::{debug, info, warn};
+use wakey_agent::protocol::{ErrorPayload, RequestId, ServerMessage};
+
+use crate::runtime::{AgentReply, AppState};
+
+#[derive(Debug, Deserialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+enum IncomingClientMessage {
+    Hello {
+        agent_id: String,
+    },
+    Auth {
+        agent_id: String,
+        agent_token: String,
+    },
+    Heartbeat {
+        agent_id: String,
+    },
+    Result {
+        request_id: RequestId,
+        result: serde_json::Value,
+    },
+    Error {
+        request_id: RequestId,
+        error: ErrorPayload,
+    },
+}
+
+pub async fn agent_ws(ws: WebSocketUpgrade, State(state): State<AppState>) -> impl IntoResponse {
+    ws.on_upgrade(move |socket| handle_agent_socket(state, socket))
+}
+
+async fn handle_agent_socket(state: AppState, socket: WebSocket) {
+    let (mut write, mut read) = socket.split();
+    let (tx, mut rx) = mpsc::unbounded_channel::<ServerMessage>();
+
+    let writer = tokio::spawn(async move {
+        while let Some(msg) = rx.recv().await {
+            let encoded = match serde_json::to_string(&msg) {
+                Ok(s) => s,
+                Err(err) => {
+                    warn!(error = %err, "failed to encode server websocket message");
+                    continue;
+                }
+            };
+            if let Err(err) = write.send(Message::Text(encoded.into())).await {
+                warn!(error = %err, "failed to send websocket message");
+                break;
+            }
+        }
+    });
+
+    let mut authed_agent_id: Option<String> = None;
+
+    loop {
+        let frame = read.next().await;
+        let msg = match frame {
+            Some(Ok(msg)) => msg,
+            Some(Err(err)) => {
+                warn!(error = %err, "agent websocket receive error");
+                break;
+            }
+            None => break,
+        };
+
+        match msg {
+            Message::Text(text) => {
+                if let Err(err) = process_agent_text(&state, &tx, &mut authed_agent_id, &text).await
+                {
+                    warn!(error = %err, "closing agent websocket due to protocol/auth error");
+                    break;
+                }
+            }
+            Message::Ping(_) => {}
+            Message::Pong(_) => {}
+            Message::Close(_) => break,
+            Message::Binary(_) => {
+                debug!("ignoring unexpected binary websocket frame");
+            }
+        }
+    }
+
+    if let Some(agent_id) = authed_agent_id {
+        info!(agent_id = %agent_id, "agent disconnected");
+        state.sessions.write().await.remove(&agent_id);
+    }
+
+    writer.abort();
+}
+
+async fn process_agent_text(
+    state: &AppState,
+    tx: &mpsc::UnboundedSender<ServerMessage>,
+    authed_agent_id: &mut Option<String>,
+    text: &str,
+) -> Result<()> {
+    let message: IncomingClientMessage =
+        serde_json::from_str(text).context("invalid client websocket payload")?;
+
+    match message {
+        IncomingClientMessage::Hello { agent_id } => {
+            debug!(agent_id = %agent_id, "agent hello received");
+        }
+        IncomingClientMessage::Auth {
+            agent_id,
+            agent_token,
+        } => {
+            if !state
+                .store
+                .verify_agent_token(&agent_id, &agent_token)
+                .await
+            {
+                anyhow::bail!("agent auth rejected");
+            }
+            state
+                .sessions
+                .write()
+                .await
+                .insert(agent_id.clone(), tx.clone());
+            *authed_agent_id = Some(agent_id.clone());
+            info!(agent_id = %agent_id, "agent authenticated");
+        }
+        IncomingClientMessage::Heartbeat { agent_id } => {
+            if authed_agent_id.as_deref() != Some(agent_id.as_str()) {
+                anyhow::bail!("heartbeat for unauthenticated or mismatched agent");
+            }
+            debug!(agent_id = %agent_id, "heartbeat received");
+        }
+        IncomingClientMessage::Result { request_id, result } => {
+            if authed_agent_id.is_none() {
+                anyhow::bail!("result before auth");
+            }
+            let key = request_id.as_str().to_string();
+            if let Some(waiter) = state.pending.lock().await.remove(&key) {
+                let _ = waiter.send(AgentReply::Result(result));
+            }
+        }
+        IncomingClientMessage::Error { request_id, error } => {
+            if authed_agent_id.is_none() {
+                anyhow::bail!("error before auth");
+            }
+            let key = request_id.as_str().to_string();
+            if let Some(waiter) = state.pending.lock().await.remove(&key) {
+                let _ = waiter.send(AgentReply::Error(error));
+            }
+        }
+    }
+
+    Ok(())
+}