2 files changed, 35 insertions, 5 deletions

diff --git a/editors/code/src/client.ts b/editors/code/src/client.ts
index 1ba2352ee..d032b45b7 100644
--- a/editors/code/src/client.ts
+++ b/editors/code/src/client.ts
@@ -2,6 +2,7 @@ import * as lc from 'vscode-languageclient/node';
 import * as vscode from 'vscode';
 import * as ra from '../src/lsp_ext';
 import * as Is from 'vscode-languageclient/lib/common/utils/is';
+import { DocumentSemanticsTokensSignature, DocumentSemanticsTokensEditsSignature, DocumentRangeSemanticTokensSignature } from 'vscode-languageclient/lib/common/semanticTokens';
 import { assert } from './util';
 
 function renderCommand(cmd: ra.CommandLink) {
@@ -18,6 +19,13 @@ function renderHoverActions(actions: ra.CommandLinkGroup[]): vscode.MarkdownStri
     return result;
 }
 
+// Workaround for https://github.com/microsoft/vscode-languageserver-node/issues/576
+async function semanticHighlightingWorkaround<R, F extends (...args: any[]) => vscode.ProviderResult<R>>(next: F, ...args: Parameters<F>): Promise<R> {
+    const res = await next(...args);
+    if (res == null) throw new Error('busy');
+    return res;
+}
+
 export function createClient(serverPath: string, cwd: string): lc.LanguageClient {
     // '.' Is the fallback if no folder is open
     // TODO?: Workspace folders support Uri's (eg: file://test.txt).
@@ -41,6 +49,15 @@ export function createClient(serverPath: string, cwd: string): lc.LanguageClient
         diagnosticCollectionName: "rustc",
         traceOutputChannel,
         middleware: {
+            provideDocumentSemanticTokens(document: vscode.TextDocument, token: vscode.CancellationToken, next: DocumentSemanticsTokensSignature): vscode.ProviderResult<vscode.SemanticTokens> {
+                return semanticHighlightingWorkaround(next, document, token);
+            },
+            provideDocumentSemanticTokensEdits(document: vscode.TextDocument, previousResultId: string, token: vscode.CancellationToken, next: DocumentSemanticsTokensEditsSignature): vscode.ProviderResult<vscode.SemanticTokensEdits | vscode.SemanticTokens> {
+                return semanticHighlightingWorkaround(next, document, previousResultId, token);
+            },
+            provideDocumentRangeSemanticTokens(document: vscode.TextDocument, range: vscode.Range, token: vscode.CancellationToken, next: DocumentRangeSemanticTokensSignature): vscode.ProviderResult<vscode.SemanticTokens> {
+                return semanticHighlightingWorkaround(next, document, range, token);
+            },
             async provideHover(document: vscode.TextDocument, position: vscode.Position, token: vscode.CancellationToken, _next: lc.ProvideHoverSignature) {
                 return client.sendRequest(lc.HoverRequest.type, client.code2ProtocolConverter.asTextDocumentPositionParams(document, position), token).then(
                     (result) => {
diff --git a/editors/code/src/snippets.ts b/editors/code/src/snippets.ts
index 258b49982..fee736e7d 100644
--- a/editors/code/src/snippets.ts
+++ b/editors/code/src/snippets.ts
@@ -3,16 +3,29 @@ import * as vscode from 'vscode';
 import { assert } from './util';
 
 export async function applySnippetWorkspaceEdit(edit: vscode.WorkspaceEdit) {
-    assert(edit.entries().length === 1, `bad ws edit: ${JSON.stringify(edit)}`);
-    const [uri, edits] = edit.entries()[0];
+    if (edit.entries().length === 1) {
+        const [uri, edits] = edit.entries()[0];
+        const editor = await editorFromUri(uri);
+        if (editor) await applySnippetTextEdits(editor, edits);
+        return;
+    }
+    for (const [uri, edits] of edit.entries()) {
+        const editor = await editorFromUri(uri);
+        if (editor) await editor.edit((builder) => {
+            for (const indel of edits) {
+                assert(!parseSnippet(indel.newText), `bad ws edit: snippet received with multiple edits: ${JSON.stringify(edit)}`);
+                builder.replace(indel.range, indel.newText);
+            }
+        });
+    }
+}
 
+async function editorFromUri(uri: vscode.Uri): Promise<vscode.TextEditor | undefined> {
     if (vscode.window.activeTextEditor?.document.uri !== uri) {
         // `vscode.window.visibleTextEditors` only contains editors whose contents are being displayed
         await vscode.window.showTextDocument(uri, {});
     }
-    const editor = vscode.window.visibleTextEditors.find((it) => it.document.uri.toString() === uri.toString());
-    if (!editor) return;
-    await applySnippetTextEdits(editor, edits);
+    return vscode.window.visibleTextEditors.find((it) => it.document.uri.toString() === uri.toString());
 }
 
 export async function applySnippetTextEdits(editor: vscode.TextEditor, edits: vscode.TextEdit[]) {