diff --git a/js/ui/endSessionDialog.js b/js/ui/endSessionDialog.js
index a10b86d00..47963bc1e 100644
--- a/js/ui/endSessionDialog.js
+++ b/js/ui/endSessionDialog.js
@@ -408,10 +408,6 @@ EndSessionDialog.prototype = {
     },
 
     _updateButtons: function() {
-        if (this.state != ModalDialog.State.OPENING &&
-            this.state != ModalDialog.State.OPENED)
-            return;
-
         let dialogContent = DialogContent[this._type];
         let buttons = [{ action: Lang.bind(this, this.cancel),
                          label:  _("Cancel"),
@@ -521,11 +517,12 @@ EndSessionDialog.prototype = {
             this._inhibitors.push(inhibitor);
         }
 
+        this._updateButtons();
+
         if (!this.open(timestamp))
             throw new DBus.DBusError('org.gnome.Shell.ModalDialog.GrabError',
                                      "Cannot grab pointer and keyboard");
 
-        this._updateButtons();
         this._updateContent();
 
         let signalId = this.connect('opened',
diff --git a/js/ui/modalDialog.js b/js/ui/modalDialog.js
index 802f9c90f..dfebc479f 100644
--- a/js/ui/modalDialog.js
+++ b/js/ui/modalDialog.js
@@ -50,10 +50,6 @@ ModalDialog.prototype = {
                                                       coordinate: Clutter.BindCoordinate.POSITION | Clutter.BindCoordinate.SIZE });
         this._group.add_constraint(constraint);
 
-        global.focus_manager.add_group(this._group);
-        this._initialKeyFocus = this._group;
-        this._savedKeyFocus = null;
-
         this._group.connect('destroy', Lang.bind(this, this._onGroupDestroy));
 
         this._actionKeys = {};
@@ -94,6 +90,10 @@ ModalDialog.prototype = {
                                { expand:  true,
                                  x_align: St.Align.MIDDLE,
                                  y_align: St.Align.END });
+
+        global.focus_manager.add_group(this._dialogLayout);
+        this._initialKeyFocus = this._dialogLayout;
+        this._savedKeyFocus = null;
     },
 
     setButtons: function(buttons) {
@@ -195,6 +195,7 @@ ModalDialog.prototype = {
 
         this.state = State.CLOSING;
         this.popModal(timestamp);
+        this._savedKeyFocus = null;
 
         Tweener.addTween(this._group,
                          { opacity: 0,