Make TranslationLayer Native AOT-compatible

src/Microsoft.TestPlatform.CommunicationUtilities/JsonDataSerializer.Stj.cs

@@ -4,25 +4,24 @@
 #if NETCOREAPP
 using System;
 using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using System.Text.Encodings.Web;
 using System.Text.Json;
 using System.Text.Json.Serialization;
+using System.Text.Json.Serialization.Metadata;
 
 using Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.Serialization;
 using Microsoft.VisualStudio.TestPlatform.Common.DataCollection;
 using Microsoft.VisualStudio.TestPlatform.ObjectModel.Client;
-using Microsoft.VisualStudio.TestPlatform.Utilities;
 
 namespace Microsoft.VisualStudio.TestPlatform.CommunicationUtilities;
 
 public partial class JsonDataSerializer
 {
-    private static readonly bool DisableFastJson = FeatureFlag.Instance.IsSet(FeatureFlag.VSTEST_DISABLE_FASTER_JSON_SERIALIZATION);
-
     private static readonly JsonSerializerOptions PayloadOptionsV1; // payload options for version <= 1
     private static readonly JsonSerializerOptions PayloadOptionsV2; // payload options for version >= 2
-    private static readonly JsonSerializerOptions FastOptions; // options for faster json
+    private static readonly JsonSerializerOptions FastOptions; // options for fast deserialization
     private static readonly JsonSerializerOptions DefaultOptions; // generic options
 
     static JsonDataSerializer()
@@ -41,6 +40,7 @@ static JsonDataSerializer()
         DefaultOptions.Converters.Add(new TestProcessAttachDebuggerPayloadConverter());
         DefaultOptions.Converters.Add(new TestSessionInfoConverter());
         DefaultOptions.Converters.Add(new DiscoveryCriteriaConverter());
+        DefaultOptions.Converters.Add(new ExceptionConverter());
 
         // V2 options: clone DefaultOptions and add V2-specific converters
         PayloadOptionsV2 = new JsonSerializerOptions(DefaultOptions);
@@ -62,6 +62,7 @@ static JsonDataSerializer()
         FastOptions = new JsonSerializerOptions(PayloadOptionsV2);
     }
 
+    [UnconditionalSuppressMessage("AOT", "IL3050", Justification = "DefaultJsonTypeInfoResolver is only used as a fallback for non-AoT builds.")]
     private static JsonSerializerOptions CreateBaseOptions() => new()
     {
         PropertyNameCaseInsensitive = true,
@@ -74,6 +75,11 @@ static JsonDataSerializer()
         NumberHandling = JsonNumberHandling.AllowReadingFromString,
         Encoder = JavaScriptEncoder.UnsafeRelaxedJsonEscaping,
         ReferenceHandler = ReferenceHandler.IgnoreCycles,
+        // Chain the source-generated context (for NativeAOT where reflection is
+        // disabled) with the default reflection-based resolver (for types not
+        // covered by the source-gen context in non-AoT builds). Under NativeAOT,
+        // DefaultJsonTypeInfoResolver is a no-op for trimmed types.
+        TypeInfoResolver = JsonTypeInfoResolver.Combine(TestPlatformJsonContext.Default, new DefaultJsonTypeInfoResolver()),
     };
 
     private static partial (int version, string? messageType) ParseHeaderFromJson(string rawMessage)
@@ -105,7 +111,7 @@ private static partial (int version, string? messageType) ParseHeaderFromJson(st
             using var doc = JsonDocument.Parse(message.RawMessage!);
             if (doc.RootElement.TryGetProperty("Payload", out var payloadElement))
             {
-                result = JsonSerializer.Deserialize<T>(payloadElement, payloadOptions);
+                result = StjSafe.Deserialize<T>(payloadElement, payloadOptions);
             }
             else
             {
@@ -180,24 +186,20 @@ private static partial string SerializeMessageCore(string? messageType)
 
     private static partial string SerializePayloadCore(string? messageType, object? payload, int version)
     {
+        if (payload is null)
+            return string.Empty;
+
         var payloadOptions = GetPayloadOptions(version);
-        // Fast json is only equivalent to the serialization that is used for protocol version 2 and upwards (or more precisely for the paths that use PayloadOptionsV2)
-        // so when we resolved the old options we should use non-fast path.
-        if (DisableFastJson || payloadOptions == PayloadOptionsV1)
-        {
-            if (payload is null)
-                return string.Empty;
 
-            var serializedPayload = JsonSerializer.SerializeToElement(payload, payloadOptions);
+        // Serialize payload to JsonElement first using the versioned options (which have the
+        // custom converters), then embed in the envelope. This two-step approach is required
+        // for NativeAOT: serializing object? Payload directly would require STJ to resolve
+        // the runtime type polymorphically via reflection.
+        var serializedPayload = StjSafe.SerializeToElement(payload, payload.GetType(), payloadOptions);
 
-            return version > 1 ?
-                Serialize(DefaultOptions, new VersionedMessageEnvelope { MessageType = messageType, Version = version, Payload = serializedPayload }) :
-                Serialize(DefaultOptions, new MessageEnvelope { MessageType = messageType, Payload = serializedPayload });
-        }
-        else
-        {
-            return Serialize(FastOptions, new VersionedMessageForSerialization { MessageType = messageType, Version = version, Payload = payload });
-        }
+        return version > 1 ?
+            Serialize(DefaultOptions, new VersionedMessageEnvelope { MessageType = messageType, Version = version, Payload = serializedPayload }) :
+            Serialize(DefaultOptions, new MessageEnvelope { MessageType = messageType, Payload = serializedPayload });
     }
 
     private static partial string SerializeCore<T>(T data, int version)
@@ -208,7 +210,7 @@ private static partial string SerializeCore<T>(T data, int version)
 
     private static T? DeserializeObjectFast<T>(string value)
     {
-        return JsonSerializer.Deserialize<T>(value, FastOptions);
+        return StjSafe.Deserialize<T>(value, FastOptions);
     }
 
     /// <summary>
@@ -220,7 +222,7 @@ private static partial string SerializeCore<T>(T data, int version)
     /// <returns>Serialized data.</returns>
     private static string Serialize<T>(JsonSerializerOptions options, T data)
     {
-        return JsonSerializer.Serialize(data, options);
+        return StjSafe.Serialize(data, options);
     }
 
     /// <summary>
@@ -232,7 +234,7 @@ private static string Serialize<T>(JsonSerializerOptions options, T data)
     /// <returns>Deserialized data.</returns>
     private static T? Deserialize<T>(JsonSerializerOptions options, string data)
     {
-        return JsonSerializer.Deserialize<T>(data, options);
+        return StjSafe.Deserialize<T>(data, options);
     }
 
     private static JsonSerializerOptions GetPayloadOptions(int? version)
@@ -258,7 +260,7 @@ private static JsonSerializerOptions GetPayloadOptions(int? version)
     /// This grabs payload from the message, we already know version and message type.
     /// </summary>
     /// <typeparam name="T"></typeparam>
-    private class PayloadedMessage<T>
+    internal class PayloadedMessage<T>
     {
         public T? Payload { get; set; }
     }
@@ -267,31 +269,31 @@ private class PayloadedMessage<T>
     /// Serialization-only DTO for building the JSON wire format (without Version).
     /// NOT a Message — this is never returned to callers.
     /// </summary>
-    private class MessageEnvelope
+    internal class MessageEnvelope
     {
         public string? MessageType { get; set; }
 
         [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
-        public object? Payload { get; set; }
+        public JsonElement? Payload { get; set; }
     }
 
     /// <summary>
     /// Serialization-only DTO for building the JSON wire format (with Version).
     /// NOT a Message — this is never returned to callers.
     /// </summary>
-    private class VersionedMessageEnvelope
+    internal class VersionedMessageEnvelope
     {
         public int Version { get; set; }
         public string? MessageType { get; set; }
 
         [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
-        public object? Payload { get; set; }
+        public JsonElement? Payload { get; set; }
     }
 
     /// <summary>
     /// For serialization directly into string, without first converting to JsonElement, and then from JsonElement to string.
     /// </summary>
-    private class VersionedMessageForSerialization
+    internal class VersionedMessageForSerialization
     {
         /// <summary>
         /// Gets or sets the version of the message

src/Microsoft.TestPlatform.CommunicationUtilities/Microsoft.TestPlatform.CommunicationUtilities.csproj

@@ -5,6 +5,16 @@
     <TargetFrameworks>$(NetFrameworkMinimum);$(ExtensionTargetFrameworks);$(NetCoreAppMinimum)</TargetFrameworks>
     <IsTestProject>false</IsTestProject>
   </PropertyGroup>
+
+  <!-- NativeAOT / trimming compatibility analysis (net8.0+ only) -->
+  <PropertyGroup Condition="$([MSBuild]::IsTargetFrameworkCompatible('$(TargetFramework)', 'net8.0'))">
+    <IsAotCompatible>true</IsAotCompatible>
+    <EnableTrimAnalyzer>true</EnableTrimAnalyzer>
+    <EnableAotAnalyzer>true</EnableAotAnalyzer>
+    <!-- Downgrade pre-existing AOT/trim warnings from errors to warnings.
+         Our new code is AOT-clean; these are from Jsonite and legacy converters. -->
+    <WarningsNotAsErrors>$(WarningsNotAsErrors);IL2026;IL2057;IL2067;IL3050</WarningsNotAsErrors>
+  </PropertyGroup>
   <ItemGroup>
     <ProjectReference Include="..\Microsoft.TestPlatform.CoreUtilities\Microsoft.TestPlatform.CoreUtilities.csproj" />
     <ProjectReference Include="..\Microsoft.TestPlatform.ObjectModel\Microsoft.TestPlatform.ObjectModel.csproj" />

src/Microsoft.TestPlatform.CommunicationUtilities/Serialization/AfterTestRunEndResultConverter.cs

@@ -52,7 +52,7 @@ public override void Write(Utf8JsonWriter writer, AfterTestRunEndResult value, J
     {
         if (element.TryGetProperty(name, out var prop) && prop.ValueKind != JsonValueKind.Null)
         {
-            return JsonSerializer.Deserialize<T>(prop.GetRawText(), options);
+            return StjSafe.Deserialize<T>(prop.GetRawText(), options);
         }
 
         return default;
@@ -61,7 +61,7 @@ public override void Write(Utf8JsonWriter writer, AfterTestRunEndResult value, J
     private static void WriteProperty<T>(Utf8JsonWriter writer, string name, T value, JsonSerializerOptions options)
     {
         writer.WritePropertyName(name);
-        JsonSerializer.Serialize(writer, value, options);
+        StjSafe.Serialize(writer, value, options);
     }
 }
 

src/Microsoft.TestPlatform.CommunicationUtilities/Serialization/AttachmentConverters.cs

@@ -32,7 +32,7 @@ internal class AttachmentSetConverter : JsonConverter<AttachmentSet>
             {
                 if (attachment.ValueKind != JsonValueKind.Null)
                 {
-                    attachmentSet.Attachments.Add(JsonSerializer.Deserialize<UriDataAttachment>(attachment.GetRawText(), options)!);
+                    attachmentSet.Attachments.Add(StjSafe.Deserialize<UriDataAttachment>(attachment.GetRawText(), options)!);
                 }
             }
         }
@@ -46,7 +46,7 @@ public override void Write(Utf8JsonWriter writer, AttachmentSet value, JsonSeria
         writer.WriteString("Uri", value.Uri.OriginalString);
         writer.WriteString("DisplayName", value.DisplayName);
         writer.WritePropertyName("Attachments");
-        JsonSerializer.Serialize(writer, value.Attachments, options);
+        StjSafe.Serialize(writer, value.Attachments, options);
         writer.WriteEndObject();
     }
 }

src/Microsoft.TestPlatform.CommunicationUtilities/Serialization/DiscoveryCriteriaConverter.cs

@@ -75,7 +75,7 @@ public override void Write(Utf8JsonWriter writer, DiscoveryCriteria value, JsonS
     {
         if (element.TryGetProperty(name, out var prop) && prop.ValueKind != JsonValueKind.Null)
         {
-            return JsonSerializer.Deserialize<T>(prop.GetRawText(), options);
+            return StjSafe.Deserialize<T>(prop.GetRawText(), options);
         }
 
         return default;
@@ -84,7 +84,7 @@ public override void Write(Utf8JsonWriter writer, DiscoveryCriteria value, JsonS
     private static void WriteProperty<T>(Utf8JsonWriter writer, string name, T value, JsonSerializerOptions options)
     {
         writer.WritePropertyName(name);
-        JsonSerializer.Serialize(writer, value, options);
+        StjSafe.Serialize(writer, value, options);
     }
 }
 

src/Microsoft.TestPlatform.CommunicationUtilities/Serialization/ExceptionConverter.cs

@@ -0,0 +1,150 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+#if NETCOREAPP
+
+using System;
+using System.Text;
+using System.Text.Json;
+using System.Text.Json.Serialization;
+
+namespace Microsoft.VisualStudio.TestPlatform.CommunicationUtilities.Serialization;
+
+/// <summary>
+/// JSON converter for <see cref="Exception"/> that handles the read-only properties
+/// (<c>Message</c>, <c>StackTrace</c>) which STJ's source-generated metadata cannot populate
+/// via the parameterless constructor.
+/// <para>
+/// On deserialization, produces a <see cref="RemoteException"/> that preserves the original
+/// type name, message, stack trace, and inner exception. The original exception type is erased
+/// (all exceptions materialize as <see cref="RemoteException"/>) but <c>ToString()</c> renders
+/// the full original information including type name and stack trace.
+/// </para>
+/// </summary>
+internal class ExceptionConverter : JsonConverter<Exception>
+{
+    /// <inheritdoc/>
+    public override Exception? Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options)
+    {
+        if (reader.TokenType == JsonTokenType.Null)
+        {
+            return null;
+        }
+
+        using var doc = JsonDocument.ParseValue(ref reader);
+        var root = doc.RootElement;
+
+        string? className = root.TryGetProperty("ClassName", out var clsProp) && clsProp.ValueKind == JsonValueKind.String
+            ? clsProp.GetString()
+            : null;
+
+        string? message = root.TryGetProperty("Message", out var msgProp) && msgProp.ValueKind != JsonValueKind.Null
+            ? msgProp.GetString()
+            : null;
+
+        string? stackTrace = root.TryGetProperty("StackTraceString", out var stProp) && stProp.ValueKind == JsonValueKind.String
+            ? stProp.GetString()
+            : null;
+
+        Exception? innerException = null;
+        if (root.TryGetProperty("InnerException", out var innerProp) && innerProp.ValueKind != JsonValueKind.Null)
+        {
+            innerException = StjSafe.Deserialize<Exception>(innerProp.GetRawText(), options);
+        }
+
+        var exception = new RemoteException(className, message, stackTrace, innerException);
+
+        if (root.TryGetProperty("HResult", out var hresultProp) && hresultProp.ValueKind == JsonValueKind.Number)
+        {
+            exception.HResult = hresultProp.GetInt32();
+        }
+
+        if (root.TryGetProperty("Source", out var sourceProp) && sourceProp.ValueKind == JsonValueKind.String)
+        {
+            exception.Source = sourceProp.GetString();
+        }
+
+        return exception;
+    }
+
+    /// <inheritdoc/>
+    public override void Write(Utf8JsonWriter writer, Exception value, JsonSerializerOptions options)
+    {
+        writer.WriteStartObject();
+
+        // For RemoteException (already deserialized once), preserve the original ClassName.
+        writer.WriteString("ClassName", value is RemoteException remote
+            ? remote.ClassName
+            : value.GetType().FullName);
+        writer.WriteString("Message", value.Message);
+        writer.WriteString("StackTraceString", value is RemoteException re
+            ? re.RemoteStackTrace
+            : value.StackTrace);
+        writer.WriteString("Source", value.Source);
+        writer.WriteNumber("HResult", value.HResult);
+
+        writer.WritePropertyName("InnerException");
+        if (value.InnerException is null)
+        {
+            writer.WriteNullValue();
+        }
+        else
+        {
+            StjSafe.Serialize(writer, value.InnerException, options);
+        }
+
+        writer.WriteEndObject();
+    }
+}
+
+/// <summary>
+/// Exception that preserves diagnostic information from a remotely-serialized exception,
+/// including the original type name and stack trace. Since the original exception type may
+/// not be available (or may be trimmed under NativeAOT), this type acts as a carrier that
+/// faithfully reproduces the original <c>ToString()</c> output.
+/// </summary>
+internal sealed class RemoteException : Exception
+{
+    /// <summary>
+    /// Gets the fully qualified name of the original exception type (e.g.
+    /// <c>"System.InvalidOperationException"</c>).
+    /// </summary>
+    public string? ClassName { get; }
+
+    /// <summary>
+    /// Gets the original stack trace string as captured from the remote process.
+    /// </summary>
+    public string? RemoteStackTrace { get; }
+
+    public RemoteException(string? className, string? message, string? stackTrace, Exception? innerException)
+        : base(message, innerException)
+    {
+        ClassName = className;
+        RemoteStackTrace = stackTrace;
+    }
+
+    public override string ToString()
+    {
+        var sb = new StringBuilder();
+        sb.Append(ClassName ?? nameof(RemoteException));
+        if (!string.IsNullOrEmpty(Message))
+        {
+            sb.Append(": ").Append(Message);
+        }
+
+        if (InnerException is not null)
+        {
+            sb.Append(" ---> ").Append(InnerException).AppendLine()
+              .Append("   --- End of inner exception stack trace ---");
+        }
+
+        if (!string.IsNullOrEmpty(RemoteStackTrace))
+        {
+            sb.AppendLine().Append(RemoteStackTrace);
+        }
+
+        return sb.ToString();
+    }
+}
+
+#endif