Support JsonElement and JObject in completion and code actions

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Completion/CompletionListMerger.cs

@@ -5,6 +5,10 @@ using System.Collections.Generic;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
+using System.Text.Json;
+using System.Text.Json.Nodes;
+using System.Text.Json.Serialization;
+using Microsoft.CodeAnalysis.Razor.Protocol;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
 using Newtonsoft.Json.Linq;
 
@@ -107,18 +111,49 @@ internal static class CompletionListMerger
             return;
         }
 
+        // We have to be agnostic to which serialization method the delegated servers use, including
+        // the scenario where they use different ones, so we normalize the data to JObject.
+        TrySplitJsonElement(data, collector);
+        TrySplitJObject(data, collector);
+    }
+
+    private static void TrySplitJsonElement(object data, List<JObject> collector)
+    {
+        if (data is not JsonElement jsonElement)
+        {
+            return;
+        }
+
+        if (jsonElement.TryGetProperty(Data1Key, out _) || jsonElement.TryGetProperty(Data1Key.ToLowerInvariant(), out _) &&
+            jsonElement.TryGetProperty(Data2Key, out _) || jsonElement.TryGetProperty(Data2Key.ToLowerInvariant(), out _))
+        {
+            // Merged data
+            var mergedCompletionListData = jsonElement.Deserialize<MergedCompletionListData>();
+
+            if (mergedCompletionListData is null)
+            {
+                Debug.Fail("Merged completion list data is null, this should never happen.");
+                return;
+            }
+
+            Split(mergedCompletionListData.Data1, collector);
+            Split(mergedCompletionListData.Data2, collector);
+        }
+        else
+        {
+            collector.Add((JObject)JsonHelpers.TryConvertFromJsonElement(jsonElement).AssumeNotNull());
+        }
+    }
+
+    private static void TrySplitJObject(object data, List<JObject> collector)
+    {
         if (data is not JObject jobject)
         {
             return;
         }
 
-        if (!(jobject.ContainsKey(Data1Key) || jobject.ContainsKey(Data1Key.ToLowerInvariant())) ||
-            !(jobject.ContainsKey(Data2Key) || jobject.ContainsKey(Data2Key.ToLowerInvariant())))
-        {
-            // Normal, non-merged data
-            collector.Add(jobject);
-        }
-        else
+        if ((jobject.ContainsKey(Data1Key) || jobject.ContainsKey(Data1Key.ToLowerInvariant())) &&
+            (jobject.ContainsKey(Data2Key) || jobject.ContainsKey(Data2Key.ToLowerInvariant())))
         {
             // Merged data
             var mergedCompletionListData = jobject.ToObject<MergedCompletionListData>();
@@ -132,6 +167,11 @@ internal static class CompletionListMerger
             Split(mergedCompletionListData.Data1, collector);
             Split(mergedCompletionListData.Data2, collector);
         }
+        else
+        {
+            // Normal, non-merged data
+            collector.Add(jobject);
+        }
     }
 
     private static void EnsureMergeableData(VSInternalCompletionList completionListA, VSInternalCompletionList completionListB)

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Protocol/JsonHelpers.cs

@@ -0,0 +1,41 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using System.Text.Json;
+using Newtonsoft.Json.Linq;
+
+namespace Microsoft.CodeAnalysis.Razor.Protocol;
+
+internal static class JsonHelpers
+{
+    private const string s_convertedFlag = "__convertedFromJsonElement";
+
+    /// <summary>
+    /// Normalizes data from JsonElement to JObject as thats what we expect to process
+    /// </summary>
+    internal static object? TryConvertFromJsonElement(object? data)
+    {
+        if (data is JsonElement element)
+        {
+            var jObject = JObject.Parse(element.GetRawText());
+            jObject[s_convertedFlag] = true;
+            return jObject;
+        }
+
+        return data;
+    }
+
+    /// <summary>
+    /// Converts from JObject back to JsonElement, but only if the original conversion was done with <see cref="TryConvertFromJsonElement(object?)"/>
+    /// </summary>
+    internal static object? TryConvertBackToJsonElement(object? data)
+    {
+        if (data is JObject jObject &&
+           jObject.ContainsKey(s_convertedFlag))
+        {
+            return JsonDocument.Parse(jObject.ToString()).RootElement;
+        }
+
+        return data;
+    }
+}

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/LanguageClient/Endpoints/RazorCustomMessageTarget_CodeActions.cs

@@ -80,7 +80,11 @@ internal partial class RazorCustomMessageTarget
 
             if (response.Response != null)
             {
-                codeActions.AddRange(response.Response);
+                foreach (var codeAction in response.Response)
+                {
+                    codeAction.Data = JsonHelpers.TryConvertFromJsonElement(codeAction.Data);
+                    codeActions.Add(codeAction);
+                }
             }
         }
 
@@ -126,6 +130,8 @@ internal partial class RazorCustomMessageTarget
 
         var textBuffer = virtualDocumentSnapshot.Snapshot.TextBuffer;
         var codeAction = resolveCodeActionParams.CodeAction;
+        codeAction.Data = JsonHelpers.TryConvertBackToJsonElement(codeAction.Data);
+
         var requests = _requestInvoker.ReinvokeRequestOnMultipleServersAsync<CodeAction, VSInternalCodeAction?>(
             textBuffer,
             Methods.CodeActionResolveName,
@@ -138,7 +144,10 @@ internal partial class RazorCustomMessageTarget
             if (response.Response is not null)
             {
                 // Only take the first response from a resolution
-                return response.Response;
+                var resolved = response.Response;
+                resolved.Data = JsonHelpers.TryConvertFromJsonElement(resolved.Data);
+
+                return resolved;
             }
         }
 

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/LanguageClient/Endpoints/RazorCustomMessageTarget_Completion.cs

@@ -155,6 +155,9 @@ internal partial class RazorCustomMessageTarget
             AddSnippetCompletions(request, ref builder.AsRef());
             completionList.Items = builder.ToArray();
 
+            completionList.Data = JsonHelpers.TryConvertFromJsonElement(completionList.Data);
+            ConvertJsonElementToJObject(completionList);
+
             return completionList;
         }
         finally
@@ -168,6 +171,14 @@ internal partial class RazorCustomMessageTarget
         }
     }
 
+    private void ConvertJsonElementToJObject(VSInternalCompletionList completionList)
+    {
+        foreach (var item in completionList.Items)
+        {
+            item.Data = JsonHelpers.TryConvertFromJsonElement(item.Data);
+        }
+    }
+
     private static TextEdit BuildRevertedEdit(TextEdit provisionalTextEdit)
     {
         TextEdit? revertedProvisionalTextEdit;
@@ -287,6 +298,8 @@ internal partial class RazorCustomMessageTarget
 
         var completionResolveParams = request.CompletionItem;
 
+        completionResolveParams.Data = JsonHelpers.TryConvertBackToJsonElement(completionResolveParams.Data);
+
         var textBuffer = virtualDocumentSnapshot.Snapshot.TextBuffer;
         var response = await _requestInvoker.ReinvokeRequestOnServerAsync<VSInternalCompletionItem, CompletionItem?>(
             textBuffer,
@@ -294,7 +307,14 @@ internal partial class RazorCustomMessageTarget
             languageServerName,
             completionResolveParams,
             cancellationToken).ConfigureAwait(false);
-        return response?.Response;
+
+        var item = response?.Response;
+        if (item is not null)
+        {
+            item.Data = JsonHelpers.TryConvertFromJsonElement(item.Data);
+        }
+
+        return item;
     }
 
     [JsonRpcMethod(LanguageServerConstants.RazorGetFormattingOptionsEndpointName, UseSingleObjectParameterDeserialization = true)]