Make Razor tag helper icons show up correctly

src/Razor/src/Microsoft.VisualStudio.LanguageServerClient.Razor/CompletionItemExtensions.cs

@@ -0,0 +1,64 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Generic;
+using Microsoft.CodeAnalysis.Razor.Completion;
+using Microsoft.VisualStudio.LanguageServer.Protocol;
+using Newtonsoft.Json.Linq;
+
+namespace Microsoft.VisualStudio.LanguageServerClient.Razor
+{
+    internal static class CompletionItemExtensions
+    {
+        private const string TagHelperElementDataKey = "_TagHelperElementData_";
+        private const string TagHelperAttributeDataKey = "_TagHelperAttributes_";
+        private const string RazorCompletionItemKind = "_CompletionItemKind_";
+
+        public static bool TryGetRazorCompletionKind(this CompletionItem completion, out RazorCompletionItemKind completionItemKind)
+        {
+            if (completion is null)
+            {
+                throw new ArgumentNullException(nameof(completion));
+            }
+
+            if (completion.Data is JObject dataObject &&
+                dataObject.TryGetValue("data", out var dataToken) &&
+                dataToken is JObject data && 
+                data.ContainsKey(RazorCompletionItemKind))
+            {
+                completionItemKind = data[RazorCompletionItemKind].ToObject<RazorCompletionItemKind>();
+                return true;
+            }
+
+            completionItemKind = default;
+            return false;
+        }
+
+        public static bool IsTagHelperElementCompletion(this CompletionItem completion)
+        {
+            if (completion.Data is JObject dataObject &&
+                dataObject.TryGetValue("data", out var dataToken) &&
+                dataToken is JObject data &&
+                data.ContainsKey(TagHelperElementDataKey))
+            {
+                return true;
+            }
+
+            return false;
+        }
+
+        public static bool IsTagHelperAttributeCompletion(this CompletionItem completion)
+        {
+            if (completion.Data is JObject dataObject &&
+                dataObject.TryGetValue("data", out var dataToken) &&
+                dataToken is JObject data &&
+                data.ContainsKey(TagHelperAttributeDataKey))
+            {
+                return true;
+            }
+
+            return false;
+        }
+    }
+}

src/Razor/src/Microsoft.VisualStudio.LanguageServerClient.Razor/DefaultRazorLanguageClientMiddleLayer.cs

@@ -4,9 +4,14 @@
 using System;
 using System.Collections.Generic;
 using System.Composition;
+using System.Linq;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.LanguageServer.Common;
+using Microsoft.CodeAnalysis.Razor.Completion;
+using Microsoft.VisualStudio.Core.Imaging;
+using Microsoft.VisualStudio.Imaging;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
+using Microsoft.VisualStudio.Text.Adornments;
 using Microsoft.VisualStudio.Threading;
 using Newtonsoft.Json.Linq;
 using Task = System.Threading.Tasks.Task;
@@ -49,7 +54,8 @@ namespace Microsoft.VisualStudio.LanguageServerClient.Razor
 
         public override bool CanHandle(string methodName)
         {
-            return methodName == Methods.TextDocumentOnTypeFormattingName;
+            return methodName == Methods.TextDocumentOnTypeFormattingName ||
+                methodName == Methods.TextDocumentCompletionName;
         }
 
         public override Task HandleNotificationAsync(string methodName, JToken methodParam, Func<JToken, Task> sendNotification)
@@ -89,10 +95,86 @@ namespace Microsoft.VisualStudio.LanguageServerClient.Razor
                 // We would have already applied the edits and moved the cursor. Return empty.
                 return emptyResult;
             }
+            else if (methodName == Methods.TextDocumentCompletionName)
+            {
+                var response = await sendRequest(methodParam).ConfigureAwait(false);
+                var result = response?.ToObject<SumType<CompletionItem[], CompletionList>?>();
+                if (!result.HasValue)
+                {
+                    return response;
+                }
+
+                result = SetCompletionIcon(result.Value);
+
+                return JToken.FromObject(result);
+            }
             else
             {
                 return await sendRequest(methodParam).ConfigureAwait(false);
             }
         }
+
+        // Internal for testing
+        internal SumType<CompletionItem[], CompletionList>? SetCompletionIcon(SumType<CompletionItem[], CompletionList> completionResult)
+        {
+            var result = completionResult.Match<SumType<CompletionItem[], CompletionList>?>(
+                items =>
+                {
+                    var newItems = items.Select(item => SetIcon(item)).ToArray();
+                    return newItems;
+                },
+                list =>
+                {
+                    var newItems = list.Items.Select(item => SetIcon(item)).ToArray();
+                    return new CompletionList()
+                    {
+                        Items = newItems,
+                        IsIncomplete = list.IsIncomplete,
+                    };
+                },
+                () => null);
+
+            return result;
+
+            static CompletionItem SetIcon(CompletionItem item)
+            {
+                ImageElement icon = null;
+                if (item.IsTagHelperElementCompletion() || item.IsTagHelperAttributeCompletion())
+                {
+                    icon = new ImageElement(new ImageId(KnownMonikers.XMLAttribute.Guid, KnownMonikers.XMLAttribute.Id));
+                }
+                else if (item.TryGetRazorCompletionKind(out var kind) &&
+                    (kind == RazorCompletionItemKind.DirectiveAttribute ||
+                    kind == RazorCompletionItemKind.DirectiveAttributeParameter ||
+                    kind == RazorCompletionItemKind.MarkupTransition))
+                {
+                    icon = new ImageElement(new ImageId(KnownMonikers.XMLAttribute.Guid, KnownMonikers.XMLAttribute.Id));
+                }
+
+                if (icon == null)
+                {
+                    return item;
+                }
+
+                return new VSCompletionItem()
+                {
+                    Label = item.Label,
+                    Kind = item.Kind,
+                    Detail = item.Detail,
+                    Documentation = item.Documentation,
+                    Preselect = item.Preselect,
+                    SortText = item.SortText,
+                    FilterText = item.FilterText,
+                    InsertText = item.InsertText,
+                    InsertTextFormat = item.InsertTextFormat,
+                    TextEdit = item.TextEdit,
+                    AdditionalTextEdits = item.AdditionalTextEdits,
+                    CommitCharacters = item.CommitCharacters,
+                    Command = item.Command,
+                    Data = item.Data,
+                    Icon = icon,
+                };
+            }
+        }
     }
 }