diff --git a/demo/demo.js b/demo/demo.js
index 5f0cfaab..e0fca7a4 100644
--- a/demo/demo.js
+++ b/demo/demo.js
@@ -157,7 +157,7 @@ exports.launch = function(env) {
     docs.svg.addFold("fold...", new Range(1, 0, 7, 0));
 
     docs.plain.addFold("fold", new Range(0, 90, 2, 30));
-    window.s = docs.plain;
+    window.s = docs.js;
     window.e = env.editor;
     setTimeout(function() {
         env.editor.selection.addEventListener("changeCursor", function() {
diff --git a/index.html b/index.html
index 7acf756a..6a7a0ea7 100644
--- a/index.html
+++ b/index.html
@@ -14,9 +14,9 @@
       <td align="right">
         <label for="doc">Document:</label>
         <select id="doc" size="1">
+          <option value="js">JavaScript Document</option>
           <option value="plain">Text Document</option>
           <option value="svg">SVG Document</option>
-          <option value="js">JavaScript Document</option>
           <option value="html">HTML Document</option>
           <option value="css">CSS Document</option>
           <option value="coffee">CoffeeScript Document</option>
diff --git a/lib/ace/document.js b/lib/ace/document.js
index 3fd06c53..09cf0488 100644
--- a/lib/ace/document.js
+++ b/lib/ace/document.js
@@ -66,11 +66,11 @@ var Document = function(text) {
         this.remove(new Range(0, 0, len, this.getLine(len-1).length));
         this.insert({row: 0, column:0}, text);
     };
-  	
+
     this.getValue = function() {
         return this.getAllLines().join(this.getNewLineCharacter());
     };
-    
+
     this.createAnchor = function(row, column) {
         return new Anchor(this, row, column);
     };
@@ -382,6 +382,7 @@ var Document = function(text) {
     this.revertDeltas = function(deltas) {
         for (var i=deltas.length-1; i>=0; i--) {
             var delta = deltas[i];
+
             var range = Range.fromPoints(delta.range.start, delta.range.end);
 
             if (delta.action == "insertLines")
diff --git a/lib/ace/edit_session.js b/lib/ace/edit_session.js
index 0eaa721c..99ea0468 100644
--- a/lib/ace/edit_session.js
+++ b/lib/ace/edit_session.js
@@ -104,12 +104,20 @@ var EditSession = function(text, mode) {
     this.onChange = function(e) {
         var delta = e.data;
         this.$modified = true;
+
+        var removedFolds = this.$updateInternalDataOnChange(e);
         if (!this.$fromUndo && this.$undoManager && !delta.ignore) {
             this.$deltas.push(delta);
+            if (removedFolds && removedFolds.length != 0) {
+                this.$deltas.push({
+                    action: "removeFolds",
+                    folds:  removedFolds
+                });
+            }
             this.$informUndoManager.schedule();
         }
 
-        this.$updateInternalDataOnChange(e);
+
         this.bgTokenizer.start(delta.range.start.row);
         this._dispatchEvent("change", e);
     };
@@ -620,26 +628,54 @@ var EditSession = function(text, mode) {
         return this.doc.remove(range);
     };
 
+    var docActions = [
+        "insertLines",
+        "insertText",
+        "removeLines",
+        "removeText"
+    ];
+
+    function filterDocActions(deltas) {
+        return deltas.filter(function(delta) {
+           return docActions.indexOf(delta.action) != -1;
+        });
+    }
+
+    function filterFoldDeltas(deltas) {
+        return deltas.filter(function(delta) {
+           return delta.action == "removeFolds";
+        });
+    }
+
     this.undoChanges = function(deltas) {
         if (!deltas.length)
             return;
 
+        var docDeltas = filterDocActions(deltas);
+        var foldDeltas = filterFoldDeltas(deltas);
         this.$fromUndo = true;
-        this.doc.revertDeltas(deltas);
+        this.doc.revertDeltas(docDeltas);
+        foldDeltas.forEach(function(foldDelta) {
+            foldDelta.folds.forEach(function(fold) {
+                 this.addFold(fold);
+            }, this);
+        }, this);
         this.$fromUndo = false;
 
-        this.$setUndoSelection(deltas, true);
+        this.$setUndoSelection(docDeltas, true);
     },
 
     this.redoChanges = function(deltas) {
         if (!deltas.length)
             return;
 
+        var docDeltas = filterDocActions(deltas);
+        var foldDeltas = filterFoldDeltas(deltas);
         this.$fromUndo = true;
-        this.doc.applyDeltas(deltas);
+        this.doc.applyDeltas(docDeltas);
         this.$fromUndo = false;
 
-        this.$setUndoSelection(deltas, false);
+        this.$setUndoSelection(docDeltas, false);
     },
 
     this.$setUndoSelection = function(deltas, isUndo) {
@@ -879,7 +915,6 @@ var EditSession = function(text, mode) {
         };
     };
 
-    // TODO: Really want to keep this name?
     this.$updateInternalDataOnChange = function(e) {
         var useWrapMode = this.$useWrapMode;
         var len;
@@ -888,8 +923,7 @@ var EditSession = function(text, mode) {
             lastRow = e.data.range.end.row,
             start = e.data.range.start,
             end = e.data.range.end;
-
-//        console.log("onChange", action, e.data.range + "");
+        var removedFolds = null;
 
         if (action.indexOf("Lines") != -1) {
             if (action == "insertLines") {
@@ -907,8 +941,8 @@ var EditSession = function(text, mode) {
                 useWrapMode && this.$wrapData.splice(firstRow, len);
 
                 var foldLines = this.$foldData;
-                var folds = this.getFoldsInRange(e.data.range);
-                this.removeFolds(folds);
+                removedFolds = this.getFoldsInRange(e.data.range);
+                this.removeFolds(removedFolds);
 
                 var foldLine = this.getFoldLine(lastRow);
                 var idx = 0;
@@ -949,8 +983,6 @@ var EditSession = function(text, mode) {
                     var cmp = foldLine.range.compareInside(start.row, start.column)
                     // Inside of the foldLine range. Need to split stuff up.
                     if (cmp == 0) {
-                        // TODO: Handle case where inseration is inside of fold!
-
                         foldLine = foldLine.split(start.row, start.column);
                         foldLine.shiftRow(len);
                         foldLine.addRemoveChars(
@@ -979,7 +1011,9 @@ var EditSession = function(text, mode) {
             len = Math.abs(e.data.range.start.column - e.data.range.end.column);
             if (action.indexOf("remove") != -1) {
                 // Get all the folds in the change range and remove them.
-                this.removeFolds(this.getFoldsInRange(e.data.range));
+                removedFolds = this.getFoldsInRange(e.data.range);
+                this.removeFolds(removedFolds);
+
                 len = -len;
             }
             var foldLine = this.getFoldLine(firstRow);
@@ -992,9 +1026,9 @@ var EditSession = function(text, mode) {
             console.error("The length of doc.$lines and $wrapData have to be the same!");
         }
 
-        // TODO:
-        // this.$updateFoldData(firstRow, lastRow);
         useWrapMode && this.$updateWrapData(firstRow, lastRow);
+
+        return removedFolds;
     };
 
     this.$updateWrapData = function(firstRow, lastRow) {