1 files changed, 8 insertions, 11 deletions

diff --git a/src/client/util/UndoManager.ts b/src/client/util/UndoManager.ts
index 6fef9d660..2379cb2ab 100644
--- a/src/client/util/UndoManager.ts
+++ b/src/client/util/UndoManager.ts
@@ -91,7 +91,7 @@ export namespace UndoManager {
     let currentBatch: UndoBatch | undefined;
     export let batchCounter = observable.box(0);
     let undoing = false;
-    let tempEvents: UndoEvent[] | undefined = undefined;
+    export let tempEvents: UndoEvent[] | undefined = undefined;
 
     export function AddEvent(event: UndoEvent, value?: any): void {
         if (currentBatch && batchCounter.get() && !undoing) {
@@ -187,15 +187,11 @@ export namespace UndoManager {
         return false;
     });
 
-    export function RunInTempBatch<T>(fn: () => T) {
+    export function StartTempBatch() {
         tempEvents = [];
-        try {
-            const success = runInAction(fn);
-            if (!success) UndoManager.UndoTempBatch();
-            return success;
-        } finally {
-            tempEvents = undefined;
-        }
+    }
+    export function EndTempBatch<T>(success: boolean) {
+        UndoManager.UndoTempBatch(success);
     }
     //TODO Make this return the return value
     export function RunInBatch<T>(fn: () => T, batchName: string) {
@@ -206,10 +202,11 @@ export namespace UndoManager {
             batch.end();
         }
     }
-    export const UndoTempBatch = action(() => {
-        if (tempEvents) {
+    export const UndoTempBatch = action((success: any) => {
+        if (tempEvents && !success) {
             undoing = true;
             for (let i = tempEvents.length - 1; i >= 0; i--) {
+                currentBatch?.includes(tempEvents[i]) && currentBatch.splice(currentBatch.indexOf(tempEvents[i]));
                 tempEvents[i].undo();
             }
             undoing = false;