1 files changed, 12 insertions, 8 deletions

diff --git a/editors/code/src/highlighting.ts b/editors/code/src/highlighting.ts
index d383d87ef..f9d2e9d90 100644
--- a/editors/code/src/highlighting.ts
+++ b/editors/code/src/highlighting.ts
@@ -5,13 +5,12 @@ const seedrandom = seedrandom_; // https://github.com/jvandemo/generator-angular
 
 import { ColorTheme, TextMateRuleSettings } from './color_theme';
 
-import { Ctx } from './ctx';
+import { Ctx, sendRequestWithRetry } from './ctx';
 
 export function activateHighlighting(ctx: Ctx) {
     const highlighter = new Highlighter(ctx);
-
-    ctx.client.onReady().then(() => {
-        ctx.client.onNotification(
+    ctx.onDidRestart(client => {
+        client.onNotification(
             'rust-analyzer/publishDecorations',
             (params: PublishDecorationsParams) => {
                 if (!ctx.config.highlightingOn) return;
@@ -31,7 +30,7 @@ export function activateHighlighting(ctx: Ctx) {
                 highlighter.setHighlights(targetEditor, params.decorations);
             },
         );
-    });
+    })
 
     vscode.workspace.onDidChangeConfiguration(
         _ => highlighter.removeHighlights(),
@@ -42,11 +41,14 @@ export function activateHighlighting(ctx: Ctx) {
         async (editor: vscode.TextEditor | undefined) => {
             if (!editor || editor.document.languageId !== 'rust') return;
             if (!ctx.config.highlightingOn) return;
+            let client = ctx.client;
+            if (!client) return;
 
             const params: lc.TextDocumentIdentifier = {
                 uri: editor.document.uri.toString(),
             };
-            const decorations = await ctx.sendRequestWithRetry<Decoration[]>(
+            const decorations = await sendRequestWithRetry<Decoration[]>(
+                client,
                 'rust-analyzer/decorationsRequest',
                 params,
             );
@@ -105,6 +107,8 @@ class Highlighter {
     }
 
     public setHighlights(editor: vscode.TextEditor, highlights: Decoration[]) {
+        let client = this.ctx.client;
+        if (!client) return;
         // Initialize decorations if necessary
         //
         // Note: decoration objects need to be kept around so we can dispose them
@@ -137,13 +141,13 @@ class Highlighter {
                 colorfulIdents
                     .get(d.bindingHash)![0]
                     .push(
-                        this.ctx.client.protocol2CodeConverter.asRange(d.range),
+                        client.protocol2CodeConverter.asRange(d.range),
                     );
             } else {
                 byTag
                     .get(d.tag)!
                     .push(
-                        this.ctx.client.protocol2CodeConverter.asRange(d.range),
+                        client.protocol2CodeConverter.asRange(d.range),
                     );
             }
         }