Fix selecting screen on undo

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

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,
           }))
         }
       },