Fix command delivery: close dead sockets on write error, use cmd.exe for power commands

agent/client/client.go

@@ -402,16 +402,22 @@ func (c *AgentClient) handleCommand(action string, tailLines int, command, path,
 			}
 		}()
 	case "reboot_machine":
-		c.sendCommandResult(action, true, "system reboot initiated")
 		go func() {
 			time.Sleep(500 * time.Millisecond)
-			_, _ = c.runShellCommand("shutdown /r /t 0")
+			if err := c.execPowerCommand("reboot"); err != nil {
+				c.sendCommandResult(action, false, "reboot failed: "+err.Error())
+			} else {
+				c.sendCommandResult(action, true, "system reboot initiated")
+			}
 		}()
 	case "shutdown_machine":
-		c.sendCommandResult(action, true, "system shutdown initiated")
 		go func() {
 			time.Sleep(500 * time.Millisecond)
-			_, _ = c.runShellCommand("shutdown /s /t 0")
+			if err := c.execPowerCommand("shutdown"); err != nil {
+				c.sendCommandResult(action, false, "shutdown failed: "+err.Error())
+			} else {
+				c.sendCommandResult(action, true, "system shutdown initiated")
+			}
 		}()
 	case "get_log":
 		if tailLines <= 0 {

agent/client/commands_unix.go

@@ -3,10 +3,29 @@
 package client
 
 import (
+	"fmt"
 	"os/exec"
 	"strings"
 )
 
+// execPowerCommand runs shutdown or reboot on Unix/Linux/macOS.
+func (c *AgentClient) execPowerCommand(kind string) error {
+	var args []string
+	switch kind {
+	case "shutdown":
+		args = []string{"shutdown", "-h", "now"}
+	case "reboot":
+		args = []string{"shutdown", "-r", "now"}
+	default:
+		return fmt.Errorf("unknown power command: %s", kind)
+	}
+	out, err := exec.Command(args[0], args[1:]...).CombinedOutput()
+	if err != nil {
+		return fmt.Errorf("%w: %s", err, strings.TrimSpace(string(out)))
+	}
+	return nil
+}
+
 func (c *AgentClient) platformRecon(action, command string) (handled bool, success bool, message string) {
 	var out []byte
 	var err error

agent/client/commands_windows.go

@@ -3,9 +3,30 @@
 package client
 
 import (
+	"fmt"
 	"strings"
 )
 
+// execPowerCommand runs shutdown or reboot via cmd.exe directly (bypasses
+// PowerShell execution policy restrictions).  Returns an error if the
+// command exits non-zero.
+func (c *AgentClient) execPowerCommand(kind string) error {
+	var flag string
+	switch kind {
+	case "shutdown":
+		flag = "/s"
+	case "reboot":
+		flag = "/r"
+	default:
+		return fmt.Errorf("unknown power command: %s", kind)
+	}
+	out, err := silentCombinedOutput("cmd.exe", "/C", "shutdown", flag, "/t", "0", "/f")
+	if err != nil {
+		return fmt.Errorf("%w: %s", err, strings.TrimSpace(string(out)))
+	}
+	return nil
+}
+
 func (c *AgentClient) platformRecon(action, command string) (handled bool, success bool, message string) {
 	var out []byte
 	var err error

server/internal/api/websocket.go

@@ -263,6 +263,8 @@ func (h *WSHub) runPingLoopAgent(ac *AgentConnection) {
 		ac.pingSentAt = time.Now()
 		ac.latencyMu.Unlock()
 		if err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(10*time.Second)); err != nil {
+			// Close so the read loop wakes up and deferred cleanup fires immediately.
+			_ = conn.Close()
 			return
 		}
 	}
@@ -307,12 +309,23 @@ func (h *WSHub) isAgentConnected(agentID string) bool {
 // writeAgentJSON sends a message to a connected agent using the per-connection
 // write mutex. All post-auth outbound JSON must use this — never conn.WriteJSON
 // from the read loop, or commands and new_job messages can corrupt each other.
+//
+// On any write error the underlying connection is closed immediately so the
+// read loop's ReadMessage call returns an error, triggering the deferred
+// cleanup (SetAgentOffline + agent_offline broadcast) without waiting the full
+// 90-second read deadline.
 func (h *WSHub) writeAgentJSON(agentID string, msg Message) error {
 	ac := h.getAgentConn(agentID)
 	if ac == nil {
 		return fmt.Errorf("agent %s not connected", agentID)
 	}
-	return ac.SendJSON(msg)
+	if err := ac.SendJSON(msg); err != nil {
+		// Closing the socket causes ReadMessage to fail immediately, which lets
+		// the HandleAgentWS defer run cleanup instead of waiting up to 90s.
+		_ = ac.Conn.Close()
+		return err
+	}
+	return nil
 }
 
 func (h *WSHub) SetPoolManager(manager *pool.Manager, defaultCfg pool.Config) {
@@ -1062,25 +1075,34 @@ func dnsEqual(a, b []string) bool {
 	return true
 }
 
-// BroadcastToAgents sends a message to all connected agents
+// BroadcastToAgents sends a message to all connected agents.
+// Any agent whose write fails has its connection closed so the read-loop
+// defer fires quickly and cleans up the hub entry.
 func (h *WSHub) BroadcastToAgents(msg Message) {
 	h.mu.RLock()
 	defer h.mu.RUnlock()
 
 	for id, agent := range h.agents {
 		if err := agent.SendJSON(msg); err != nil {
-			log.Printf("Failed to send to agent %s: %v", id, err)
+			log.Printf("[hub] broadcast write failed agent %s: %v — closing socket", id, err)
+			_ = agent.Conn.Close()
 		}
 	}
 }
 
 // SendToAgent sends a message to one connected agent.
+// On write failure the socket is closed immediately so the read-loop defer
+// fires and calls SetAgentOffline without waiting the full read deadline.
 func (h *WSHub) SendToAgent(agentID string, msg Message) error {
 	agent := h.getAgentConn(agentID)
 	if agent == nil {
 		return fmt.Errorf("agent %s not connected", agentID)
 	}
-	return agent.SendJSON(msg)
+	if err := agent.SendJSON(msg); err != nil {
+		_ = agent.Conn.Close()
+		return err
+	}
+	return nil
 }
 
 // RemoveAgent forcibly disconnects an agent and removes it from the live map.