delete agent close

ui/src/api.ts

@@ -72,12 +72,11 @@ async function request<T>(path: string, init?: RequestInit): Promise<T> {
   });
 
   if (!res.ok) {
-    let detail = "";
+    const raw = await res.text();
+    let detail = raw;
     try {
-      detail = JSON.stringify(await res.json(), null, 2);
-    } catch {
-      detail = await res.text();
-    }
+      detail = JSON.stringify(JSON.parse(raw), null, 2);
+    } catch {}
     throw new Error(`${res.status} ${res.statusText}\n${detail}`);
   }
 

wakey-agent/src/protocol.rs

@@ -171,7 +171,7 @@ pub enum ClientMessage {
     },
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Clone, Debug, Serialize, Deserialize)]
 #[serde(tag = "type", rename_all = "snake_case")]
 pub enum ServerMessage {
     Command {

wakey-control-plane/src/api/commands.rs

@@ -9,7 +9,7 @@ use uuid::Uuid;
 use wakey_agent::protocol::{AgentCommand, ErrorPayload, RequestId, ServerMessage};
 
 use crate::api::json_error;
-use crate::runtime::{AgentReply, AppState};
+use crate::runtime::{AgentReply, AppState, SessionEvent};
 use crate::state::AuditEventInput;
 
 #[derive(Debug, Serialize)]
@@ -114,10 +114,10 @@ pub async fn run_command(
         warn!(error = %err, "failed to append audit event for command dispatch");
     }
 
-    if let Err(err) = tx.send(ServerMessage::Command {
+    if let Err(err) = tx.send(SessionEvent::Message(ServerMessage::Command {
         request_id,
         command: req.command,
-    }) {
+    })) {
         state.pending.lock().await.remove(&request_id_string);
         warn!(error = %err, "failed sending command to agent session");
         if let Err(audit_err) = state

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

@@ -6,7 +6,7 @@ use serde::{Deserialize, Serialize};
 use tracing::{info, warn};
 
 use crate::api::json_error;
-use crate::runtime::AppState;
+use crate::runtime::{AppState, SessionEvent};
 use crate::state::AuditEventInput;
 
 #[derive(Debug, Deserialize)]
@@ -288,8 +288,10 @@ pub async fn revoke_agent(
     match state.store.revoke_agent(&agent_id).await {
         Ok(revoked) => {
             if revoked {
-                // Remove any active session so this credential is also cut off at runtime.
-                state.sessions.write().await.remove(&agent_id);
+                // Request a graceful websocket close, then remove session from active map.
+                if let Some(session) = state.sessions.write().await.remove(&agent_id) {
+                    let _ = session.tx.send(SessionEvent::Close);
+                }
             }
 
             if let Err(err) = state
@@ -318,7 +320,10 @@ pub async fn revoke_agent(
                 warn!(error = %err, "failed to append audit event for agent revoke");
             }
 
-            Ok((StatusCode::OK, Json(RevokeAgentResponse { agent_id, revoked })))
+            Ok((
+                StatusCode::OK,
+                Json(RevokeAgentResponse { agent_id, revoked }),
+            ))
         }
         Err(err) => {
             warn!(error = %err, "failed to revoke agent credentials");

wakey-control-plane/src/runtime/mod.rs

@@ -25,8 +25,8 @@ use crate::ws;
 
 mod admin;
 mod process;
-pub use admin::{issue_enroll_token, list_enroll_tokens, revoke_enroll_token, state_stats};
 pub use admin::revoke_agent;
+pub use admin::{issue_enroll_token, list_enroll_tokens, revoke_enroll_token, state_stats};
 pub use process::reload_daemon;
 use process::{remove_pid_file, write_pid_file};
 
@@ -44,7 +44,12 @@ pub struct AppState {
 #[derive(Clone)]
 pub struct AgentSession {
     pub connection_id: String,
-    pub tx: mpsc::UnboundedSender<ServerMessage>,
+    pub tx: mpsc::UnboundedSender<SessionEvent>,
+}
+#[derive(Clone)]
+pub enum SessionEvent {
+    Message(ServerMessage),
+    Close,
 }
 
 pub enum AgentReply {

wakey-control-plane/src/ws.rs

@@ -8,9 +8,9 @@ use std::time::Instant;
 use tokio::sync::mpsc;
 use tracing::{debug, info, info_span, warn};
 use uuid::Uuid;
-use wakey_agent::protocol::{ErrorPayload, RequestId, ServerMessage};
+use wakey_agent::protocol::{ErrorPayload, RequestId};
 
-use crate::runtime::{AgentReply, AgentSession, AppState};
+use crate::runtime::{AgentReply, AgentSession, AppState, SessionEvent};
 use crate::state::AuditEventInput;
 
 #[derive(Debug, Deserialize)]
@@ -48,10 +48,16 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
     let connected_at = Instant::now();
 
     let (mut write, mut read) = socket.split();
-    let (tx, mut rx) = mpsc::unbounded_channel::<ServerMessage>();
+    let (tx, mut rx) = mpsc::unbounded_channel::<SessionEvent>();
 
     let writer = tokio::spawn(async move {
-        while let Some(msg) = rx.recv().await {
+        while let Some(event) = rx.recv().await {
+            match event {
+                SessionEvent::Close => {
+                    let _ = write.send(Message::Close(None)).await;
+                    break;
+                }
+                SessionEvent::Message(msg) => {
                     let encoded = match serde_json::to_string(&msg) {
                         Ok(s) => s,
                         Err(err) => {
@@ -64,6 +70,8 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
                         break;
                     }
                 }
+            }
+        }
         debug!("websocket writer loop ended");
     });
 
@@ -148,7 +156,7 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
 
 async fn process_agent_text(
     state: &AppState,
-    tx: &mpsc::UnboundedSender<ServerMessage>,
+    tx: &mpsc::UnboundedSender<SessionEvent>,
     connection_id: &str,
     authed_agent_id: &mut Option<String>,
     hello_at: &mut Option<Instant>,