diff --git a/packages/builder/src/stores/builder/automations.ts b/packages/builder/src/stores/builder/automations.ts
index 9b20b4cd03..98213581a6 100644
--- a/packages/builder/src/stores/builder/automations.ts
+++ b/packages/builder/src/stores/builder/automations.ts
@@ -1437,8 +1437,6 @@ class AutomationStore extends BudiStore<AutomationState> {
     this.history = createHistoryStore({
       getDoc: this.actions.getDefinition.bind(this),
       selectDoc: this.actions.select.bind(this),
-      beforeAction: () => {},
-      afterAction: () => {},
     })
 
     // Then wrap save and delete with history
diff --git a/packages/builder/src/stores/builder/screens.ts b/packages/builder/src/stores/builder/screens.ts
index 5163c6a3ea..a1b51875b3 100644
--- a/packages/builder/src/stores/builder/screens.ts
+++ b/packages/builder/src/stores/builder/screens.ts
@@ -58,13 +58,12 @@ export class ScreenStore extends BudiStore<ScreenState> {
       getDoc: (id: string) =>
         get(this.store).screens?.find(screen => screen._id === id),
       selectDoc: this.select,
-      beforeAction: () => {},
       afterAction: () => {
         // Ensure a valid component is selected
         if (!get(selectedComponent)) {
-          this.update(state => ({
+          componentStore.update(state => ({
             ...state,
-            selectedComponentId: get(selectedScreen)?.props._id,
+            selectedComponentId: get(selectedScreen)?._id,
           }))
         }
       },