feat(ws): handle shutdown message for graceful destroy

protocol/messages.go

@@ -25,6 +25,11 @@ const (
 	MessageTypeTaskRequest MessageType = "task.request"
 	MessageTypeTaskCancel  MessageType = "task.cancel"
 	MessageTypeMCPCall     MessageType = "mcp.call"
+	// Shutdown is Safari telling the runner to drop everything and exit —
+	// used when CloudManager tears down a cloud sandbox so the runner's
+	// WebSocket closes immediately instead of waiting for the 5s force
+	// unregister timeout.
+	MessageTypeShutdown MessageType = "shutdown"
 )
 
 // Message is the base WebSocket message structure.

ws/handler.go

@@ -80,6 +80,8 @@ func (h *Handler) Handle(ctx context.Context, msg *protocol.Message) error {
 		return h.handleTaskCancel(ctx, msg)
 	case protocol.MessageTypeMCPCall:
 		return h.handleMCPCall(ctx, msg)
+	case protocol.MessageTypeShutdown:
+		return h.handleShutdown()
 	default:
 		slog.Warn("unknown message type",
 			"type", msg.Type,
@@ -88,6 +90,23 @@ func (h *Handler) Handle(ctx context.Context, msg *protocol.Message) error {
 	}
 }
 
+// handleShutdown cancels running tasks and closes the WebSocket client.
+// Triggered by Safari when it's tearing down the sandbox (cloud sandbox
+// recycle, eviction, explicit destroy). Close() sets c.closed so
+// RunWithReconnect returns nil instead of attempting to reconnect.
+func (h *Handler) handleShutdown() error {
+	slog.Info("shutdown message received, cancelling tasks and closing connection")
+	h.CancelAllTasks()
+	// Brief window for in-flight tool results to flush before the socket
+	// closes. Safari's CloudManager.disconnectRunner waits 5s before forcing
+	// the unregister — stay well inside that.
+	_ = h.WaitForTasks(2 * time.Second)
+	if h.client != nil {
+		_ = h.client.Close()
+	}
+	return nil
+}
+
 func (h *Handler) handleTaskRequest(ctx context.Context, msg *protocol.Message) error {
 	var req protocol.TaskRequestPayload
 	if err := json.Unmarshal(msg.Payload, &req); err != nil {