diff --git a/code/GridFieldPageSectionsExtension.php b/code/GridFieldPageSectionsExtension.php
index c4d0649..922cdb8 100644
--- a/code/GridFieldPageSectionsExtension.php
+++ b/code/GridFieldPageSectionsExtension.php
@@ -20,9 +20,9 @@ class GridFieldPageSectionsExtension implements
 	);
 
 
-	public function __construct($page, $allowedElementsOnPage = array(), $sortField = "SortOrder") {
+	public function __construct($page, $allowedParentElements = array(), $sortField = "SortOrder") {
 		$this->page = $page;
-		$this->allowedElementsOnPage = $allowedElementsOnPage;
+		$this->allowedParentElements = $allowedParentElements;
 		$this->sortField = $sortField;
 	}
 
@@ -130,20 +130,20 @@ public function getColumnAttributes($gridField, $record, $columnName) {
 			// if element has no parent we need to
 			// know the allowed elements of the page
 			if (!$record->_Parent) {
-				$pageClasses = $this->allowedElementsOnPage;
-				$pageElems = array();
-				foreach ($pageClasses as $class) {
-					$pageElems[$class] = $class::$singular_name;
+				$parentClasses = $this->allowedParentElements;
+				$parentElems = array();
+				foreach ($parentClasses as $class) {
+					$parentElems[$class] = $class::$singular_name;
 				}
 			}
 
 			return array(
-				"class"                      => "col-treenav",
-				"data-class"                 => $record->ClassName,
-				"data-level"                 => strval($record->_Level),
-				"data-parent"                => $record->_Parent ? strval($record->_Parent->ID) : "",
-				"data-allowed-page-elements" => !$record->_Parent ? json_encode($pageElems, JSON_UNESCAPED_UNICODE) : "",
-				"data-allowed-elements"      => json_encode($elems, JSON_UNESCAPED_UNICODE),
+				"class"                        => "col-treenav",
+				"data-class"                   => $record->ClassName,
+				"data-level"                   => strval($record->_Level),
+				"data-parent"                  => $record->_Parent ? strval($record->_Parent->ID) : "",
+				"data-allowed-parent-elements" => !$record->_Parent ? json_encode($parentElems, JSON_UNESCAPED_UNICODE) : "",
+				"data-allowed-elements"        => json_encode($elems, JSON_UNESCAPED_UNICODE),
 			);
 		}
 
diff --git a/code/PageElement.php b/code/PageElement.php
index ec05aa5..09d7a24 100644
--- a/code/PageElement.php
+++ b/code/PageElement.php
@@ -77,7 +77,7 @@ public function getChildrenGridField() {
 			->addComponent($dataColumns = new GridFieldDataColumns())
 			->addComponent($autoCompl)
 			->addComponent($addNewButton)
-			->addComponent(new GridFieldPageSectionsExtension($this->owner))
+			->addComponent(new GridFieldPageSectionsExtension($this->owner, $this->owner->getAllowedPageElements()))
 			->addComponent(new GridFieldDetailForm())
 			->addComponent(new GridFieldFooter());
 		$dataColumns->setFieldCasting(array('GridFieldPreview' => 'HTMLText->RAW'));
diff --git a/javascript/GridFieldPageSectionsExtension.js b/javascript/GridFieldPageSectionsExtension.js
index d1b2a61..1b19ec6 100644
--- a/javascript/GridFieldPageSectionsExtension.js
+++ b/javascript/GridFieldPageSectionsExtension.js
@@ -170,11 +170,6 @@
 								var newParent = type === "child" ? $this.data("id") : $treenav.data("parent");
 								var sort = type === "child" ? childOrder : $reorder.data("sort");
 
-								// we alter the state of the published / saved buttons
-								$('.cms-edit-form .Actions #Form_EditForm_action_publish').button({
-									showingAlternate: true
-								});
-
 								grid.reload({
 									url: grid.data("url-reorder"),
 									data: [{
@@ -194,6 +189,12 @@
 										value: sort,
 									}],
 								});
+								// we alter the state of the published / saved buttons
+								$('.cms-edit-form .Actions #Form_EditForm_action_publish').button({
+									showingAlternate: true
+								});
+								$('.cms-preview').entwine('.ss.preview').changeState('StageLink');
+
 							},
 						});
 					});
@@ -250,8 +251,7 @@
 										)
 									)
 								) {
-									var allowed = $treenav.data("allowed-page-elements");
-									console.log(allowed, element, $treenav)
+									var allowed = $treenav.data("allowed-parent-elements");
 									if (!allowed[element]) return;
 								}
 								// check our allowed children, or the allowed children of our parent row.