1 files changed, 19 insertions, 3 deletions
diff --git a/src/client/util/UndoManager.ts b/src/client/util/UndoManager.ts
index eb13ff1ee..10aa6d523 100644
--- a/src/client/util/UndoManager.ts
+++ b/src/client/util/UndoManager.ts
@@ -35,7 +35,20 @@ function propertyDecorator(target: any, key: string | symbol) {
         }
     })
 }
-export function undoBatch(target: any, key: string | symbol, descriptor?: TypedPropertyDescriptor<any>): any {
+
+export function undoBatch(target: any, key: string | symbol, descriptor?: TypedPropertyDescriptor<any>): any;
+export function undoBatch(fn: (...args: any[]) => any): (...args: any[]) => any;
+export function undoBatch(target: any, key?: string | symbol, descriptor?: TypedPropertyDescriptor<any>): any {
+    if (!key) {
+        return function () {
+            let batch = UndoManager.StartBatch("");
+            try {
+                return target.apply(undefined, arguments)
+            } finally {
+                batch.end();
+            }
+        }
+    }
     if (!descriptor) {
         propertyDecorator(target, key);
         return;
@@ -129,8 +142,11 @@ export namespace UndoManager {
 
     export function RunInBatch(fn: () => void, batchName: string) {
         let batch = StartBatch(batchName);
-        fn();
-        batch.end();
+        try {
+            fn();
+        } finally {
+            batch.end();
+        }
     }
 
     export const Undo = action(() => {