diff --git a/tui/confirm_view.go b/tui/confirm_view.go
index 7925d24..8226f4c 100644
--- a/tui/confirm_view.go
+++ b/tui/confirm_view.go
@@ -57,9 +57,6 @@ func (v *ConfirmView) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	switch msg := msg.(type) {
 	case tea.WindowSizeMsg:
 		v.width, v.height = msg.Width, msg.Height
-	case ConfirmMessage:
-		v.tipLabel.Text = msg.Message
-		v.confirmButton.OnClick = msg.ConfirmAction
 	case tea.KeyMsg:
 
 	}
@@ -68,3 +65,9 @@ func (v *ConfirmView) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	cmds = append(cmds, cmd)
 	return v, tea.Batch(cmds...)
 }
+
+func (v *ConfirmView) Reset(tip string, confirmOnClick, cancelOnClick func() tea.Cmd) {
+	v.tipLabel.Text = tip
+	v.confirmButton.OnClick = confirmOnClick
+	v.cancelButton.OnClick = cancelOnClick
+}
diff --git a/tui/editor_view.go b/tui/editor_view.go
index efba002..c0b8fac 100644
--- a/tui/editor_view.go
+++ b/tui/editor_view.go
@@ -73,9 +73,6 @@ func (v *EditorView) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	case AppliedNewHostContent:
 		v.SetHostNode(v.hostNode)
 	case tea.KeyMsg:
-		if v.model.state != StateEditorView {
-			return v, nil
-		}
 		switch {
 		case key.Matches(m, keys.Esc):
 			return v, nil
diff --git a/tui/main_view.go b/tui/main_view.go
index 0459bb2..e165ac4 100644
--- a/tui/main_view.go
+++ b/tui/main_view.go
@@ -102,13 +102,25 @@ func (m *Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 				m.switchState(m.preState)
 				m.helpView.helpView.ShowAll = false
 			}
+		default:
+			switch m.state {
+			case StateTreeView:
+				m.updateView(msg, &cmds, m.treeView)
+			case StateEditorView:
+				m.updateView(msg, &cmds, m.editorView)
+			case StateNodeView:
+				m.updateView(msg, &cmds, m.nodeView)
+			case StateConfirmView:
+				m.updateView(msg, &cmds, m.confirmView)
+			}
 		}
+	default:
+		m.updateView(msg, &cmds, m.editorView)
+		m.updateView(msg, &cmds, m.nodeView)
+		m.updateView(msg, &cmds, m.treeView)
+		m.updateView(msg, &cmds, m.confirmView)
+		m.updateView(msg, &cmds, m.helpView)
 	}
-	m.updateView(msg, &cmds, m.confirmView)
-	m.updateView(msg, &cmds, m.editorView)
-	m.updateView(msg, &cmds, m.nodeView)
-	m.updateView(msg, &cmds, m.treeView)
-	m.updateView(msg, &cmds, m.helpView)
 	return m, tea.Batch(cmds...)
 }
 
diff --git a/tui/msg.go b/tui/msg.go
index 4f202c3..498502e 100644
--- a/tui/msg.go
+++ b/tui/msg.go
@@ -1,14 +1,8 @@
 package tui
 
-import tea "github.com/charmbracelet/bubbletea"
-
 type RefreshTreeViewItems struct {
 }
 
 type AppliedNewHostContent struct {
 }
 
-type ConfirmMessage struct {
-	Message       string
-	ConfirmAction func() tea.Cmd
-}
diff --git a/tui/node_view.go b/tui/node_view.go
index 47f7631..e9bba7e 100644
--- a/tui/node_view.go
+++ b/tui/node_view.go
@@ -50,17 +50,17 @@ func NewNodeView(model *Model) *NodeView {
 	confirmButton.OnClick = func() tea.Cmd {
 		// Check inputs
 		if nodeTypeChoices.SelectedItem() == nil {
-			model.switchState(StateConfirmView)
-			return func() tea.Msg {
-				return ConfirmMessage{
-					Message: "Please select node type",
-					ConfirmAction: func() tea.Cmd {
-						// Go back to previous state
-						model.setState(model.preState)
-						return nil
-					},
-				}
-			}
+			model.confirmView.Reset("Please select node type",
+				func() tea.Cmd {
+					model.setState(StateNodeView)
+					return nil
+				},
+				func() tea.Cmd {
+					model.setState(StateNodeView)
+					return nil
+				})
+			model.setState(StateConfirmView)
+			return nil
 		}
 
 		// Get parent group
@@ -136,9 +136,6 @@ func (v *NodeView) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	case tea.WindowSizeMsg:
 		log.Debug(fmt.Sprintf("node view w %d h %d", m.Width, m.Height))
 	case tea.KeyMsg:
-		if v.model.state != StateNodeView {
-			return v, nil
-		}
 		return v.Form.Update(msg)
 	}
 	_, cmd := v.Form.Update(msg)
diff --git a/tui/tree_view.go b/tui/tree_view.go
index 427438e..3678f72 100644
--- a/tui/tree_view.go
+++ b/tui/tree_view.go
@@ -149,9 +149,6 @@ func (v *TreeView) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	case RefreshTreeViewItems, AppliedNewHostContent:
 		v.RefreshTreeNodes()
 	case tea.KeyMsg:
-		if v.model.state != StateTreeView {
-			return v, nil
-		}
 		selectedNode := v.SelectedNode()
 		switch {
 		case key.Matches(m, keys.Esc):