diff --git a/lib/ace/edit_session.js b/lib/ace/edit_session.js
index b563b534..3f1810cc 100644
--- a/lib/ace/edit_session.js
+++ b/lib/ace/edit_session.js
@@ -808,7 +808,8 @@ var EditSession = function(text, mode) {
         if (!this.$modes[textModeId])
             this.$modes[textModeId] = new TextMode();
 
-        var tokenizer = this.$modes[textModeId].getTokenizer();
+        this.$mode = this.$modes[textModeId];
+        var tokenizer = this.$mode.getTokenizer();
 
         if (!this.bgTokenizer) {
             this.bgTokenizer = new BackgroundTokenizer(tokenizer);
@@ -819,9 +820,10 @@ var EditSession = function(text, mode) {
         } else {
             this.bgTokenizer.setTokenizer(tokenizer);
         }
+        this.bgTokenizer.setDocument(this.getDocument());
 
-        this.tokenRe = this.$modes[textModeId].tokenRe;
-        this.nonTokenRe = this.$modes[textModeId].nonTokenRe;
+        this.tokenRe = this.$mode.tokenRe;
+        this.nonTokenRe = this.$mode.nonTokenRe;
     };
 
     /**