Fix the TranscribeOnly bug (take two) (#2026)

We were using our own UI, not the byzantine internal UI where it
actually needed to be fixed. Whole lot of reflection.

Also had to fix our `CoreCLR` compiler constant.

Co-authored-by: Patrick Meinecke <SeeminglyScience@users.noreply.github.com>

Author: andyleejordan

.editorconfig

@@ -205,6 +205,8 @@ dotnet_diagnostic.IDE0052.severity = error
 dotnet_diagnostic.IDE0053.severity = error
 # IDE0054: Use compound assignment
 dotnet_diagnostic.IDE0054.severity = error
+# IDE0059: Unnecessary assignment of a value
+dotnet_diagnostic.IDE0059.severity = error
 # IDE0063: Use simple 'using' statement
 dotnet_diagnostic.IDE0063.severity = error
 # IDE0066: Use switch expression

src/PowerShellEditorServices.Hosting/PowerShellEditorServices.Hosting.csproj

@@ -6,7 +6,7 @@
     <AssemblyName>Microsoft.PowerShell.EditorServices.Hosting</AssemblyName>
   </PropertyGroup>
 
-  <PropertyGroup Condition=" '$(TargetFramework)' == 'net6.0' ">
+  <PropertyGroup Condition=" '$(TargetFramework)' != 'net462' ">
     <DefineConstants>$(DefineConstants);CoreCLR</DefineConstants>
   </PropertyGroup>
 

src/PowerShellEditorServices.VSCode/PowerShellEditorServices.VSCode.csproj

@@ -6,7 +6,11 @@
     <Description>Provides added functionality to PowerShell Editor Services for the Visual Studio Code editor.</Description>
     <TargetFrameworks>netstandard2.0</TargetFrameworks>
     <AssemblyName>Microsoft.PowerShell.EditorServices.VSCode</AssemblyName>
-    <Configurations>Debug;Release;CoreCLR</Configurations>
+    <Configurations>Debug;Release</Configurations>
+  </PropertyGroup>
+
+  <PropertyGroup Condition=" '$(TargetFramework)' != 'net462' ">
+    <DefineConstants>$(DefineConstants);CoreCLR</DefineConstants>
   </PropertyGroup>
 
   <!-- Fail the release build if there are missing public API documentation comments -->

src/PowerShellEditorServices/PowerShellEditorServices.csproj

@@ -9,6 +9,10 @@
     <Configurations>Debug;Release</Configurations>
   </PropertyGroup>
 
+  <PropertyGroup Condition=" '$(TargetFramework)' != 'net462' ">
+    <DefineConstants>$(DefineConstants);CoreCLR</DefineConstants>
+  </PropertyGroup>
+
   <ItemGroup>
     <AssemblyAttribute Include="System.Runtime.CompilerServices.InternalsVisibleTo">
       <_Parameter1>Microsoft.PowerShell.EditorServices.Hosting</_Parameter1>

src/PowerShellEditorServices/Services/PowerShell/Execution/SynchronousPowerShellTask.cs

@@ -105,6 +105,12 @@ private IReadOnlyList<TResult> ExecuteNormally(CancellationToken cancellationTok
             if (PowerShellExecutionOptions.WriteOutputToHost)
             {
                 _psCommand.AddOutputCommand();
+
+                // Fix the transcription bug!
+                if (!_pwsh.Runspace.RunspaceIsRemote)
+                {
+                    _psesHost.DisableTranscribeOnly();
+                }
             }
 
             cancellationToken.Register(CancelNormalExecution);
@@ -148,7 +154,7 @@ private IReadOnlyList<TResult> ExecuteNormally(CancellationToken cancellationTok
                 if (e is PSRemotingTransportException)
                 {
                     _ = System.Threading.Tasks.Task.Run(
-                        () => _psesHost.UnwindCallStack(),
+                        _psesHost.UnwindCallStack,
                         CancellationToken.None)
                         .HandleErrorsAsync(_logger);
 
@@ -189,8 +195,6 @@ private IReadOnlyList<TResult> ExecuteNormally(CancellationToken cancellationTok
 
         private IReadOnlyList<TResult> ExecuteInDebugger(CancellationToken cancellationToken)
         {
-            // TODO: How much of this method can we remove now that it only processes PowerShell's
-            // intrinsic debugger commands?
             cancellationToken.Register(CancelDebugExecution);
 
             PSDataCollection<PSObject> outputCollection = new();
@@ -247,7 +251,7 @@ private IReadOnlyList<TResult> ExecuteInDebugger(CancellationToken cancellationT
                 if (e is PSRemotingTransportException)
                 {
                     _ = System.Threading.Tasks.Task.Run(
-                        () => _psesHost.UnwindCallStack(),
+                        _psesHost.UnwindCallStack,
                         CancellationToken.None)
                         .HandleErrorsAsync(_logger);
 

src/PowerShellEditorServices/Services/PowerShell/Host/EditorServicesConsolePSHostUserInterface.cs

@@ -7,39 +7,21 @@
 using System.Collections.ObjectModel;
 using System.Management.Automation;
 using System.Management.Automation.Host;
-using System.Reflection;
 using System.Security;
 using Microsoft.Extensions.Logging;
-using Microsoft.PowerShell.EditorServices.Utility;
 
 namespace Microsoft.PowerShell.EditorServices.Services.PowerShell.Host
 {
     internal class EditorServicesConsolePSHostUserInterface : PSHostUserInterface, IHostUISupportsMultipleChoiceSelection
     {
         private readonly PSHostUserInterface _underlyingHostUI;
 
-        private static readonly Action<PSHostUserInterface, bool> s_setTranscribeOnlyDelegate;
-
         /// <summary>
         /// We use a ConcurrentDictionary because ConcurrentHashSet does not exist, hence the value
         /// is never actually used, and `WriteProgress` must be thread-safe.
         /// </summary>
         private readonly ConcurrentDictionary<(long, int), object> _currentProgressRecords = new();
 
-        static EditorServicesConsolePSHostUserInterface()
-        {
-            if (VersionUtils.IsPS5)
-            {
-                PropertyInfo transcribeOnlyProperty = typeof(PSHostUserInterface)
-                    .GetProperty("TranscribeOnly", BindingFlags.NonPublic | BindingFlags.Instance);
-
-                MethodInfo transcribeOnlySetMethod = transcribeOnlyProperty.GetSetMethod(nonPublic: true);
-
-                s_setTranscribeOnlyDelegate = (Action<PSHostUserInterface, bool>)Delegate.CreateDelegate(
-                    typeof(Action<PSHostUserInterface, bool>), transcribeOnlySetMethod);
-            }
-        }
-
         public EditorServicesConsolePSHostUserInterface(
             ILoggerFactory loggerFactory,
             PSHostUserInterface underlyingHostUI)
@@ -105,17 +87,6 @@ internal void ResetProgress()
             // TODO: Maybe send the OSC sequence to turn off progress indicator.
         }
 
-        // This works around a bug in PowerShell 5.1 (that was later fixed) where a running
-        // transcription could cause output to disappear since the `TranscribeOnly` property was
-        // accidentally not reset to false.
-        internal void DisableTranscribeOnly()
-        {
-            if (VersionUtils.IsPS5)
-            {
-                s_setTranscribeOnlyDelegate(_underlyingHostUI, false);
-            }
-        }
-
         public override void WriteVerboseLine(string message) => _underlyingHostUI.WriteVerboseLine(message);
 
         public override void WriteWarningLine(string message) => _underlyingHostUI.WriteWarningLine(message);

src/PowerShellEditorServices/Services/PowerShell/Host/PsesInternalHost.cs

@@ -39,6 +39,31 @@ internal class PsesInternalHost : PSHost, IHostSupportsInteractiveSession, IRuns
 
         private static readonly PropertyInfo s_scriptDebuggerTriggerObjectProperty;
 
+#if !CoreCLR
+        /// <summary>
+        /// To workaround a horrid bug where the `TranscribeOnly` field of the PSHostUserInterface
+        /// can accidentally remain true, we have to use a bunch of reflection so that <see
+        /// cref="DisableTranscribeOnly()" /> can reset it to false. (This was fixed in PowerShell
+        /// 7.) Note that it must be the internal UI instance, not our own UI instance, otherwise
+        /// this would be easier. Because of the amount of reflection involved, we contain it to
+        /// only PowerShell 5.1 at compile-time, and we have to set this up in this class, not <see
+        /// cref="SynchronousPowerShellTask" /> because that's templated, making statics practically
+        /// useless. <see cref = "SynchronousPowerShellTask.ExecuteNormally" /> method calls <see
+        /// cref="DisableTranscribeOnly()" /> when necessary.
+        /// See: https://github.com/PowerShell/PowerShell/pull/3436
+        /// </summary>
+        [ThreadStatic] // Because we can re-use it, but only once per instance of PSES.
+        private static PSHostUserInterface s_internalPSHostUserInterface;
+
+        private static readonly Func<PSHostUserInterface, bool> s_getTranscribeOnlyDelegate;
+
+        private static readonly Action<PSHostUserInterface, bool> s_setTranscribeOnlyDelegate;
+
+        private static readonly PropertyInfo s_executionContextProperty;
+
+        private static readonly PropertyInfo s_internalHostProperty;
+#endif
+
         private readonly ILoggerFactory _loggerFactory;
 
         private readonly ILogger _logger;
@@ -104,6 +129,27 @@ static PsesInternalHost()
             s_scriptDebuggerTriggerObjectProperty = scriptDebuggerType.GetProperty(
                 "TriggerObject",
                 BindingFlags.Instance | BindingFlags.NonPublic);
+
+#if !CoreCLR
+            PropertyInfo transcribeOnlyProperty = typeof(PSHostUserInterface)
+                .GetProperty("TranscribeOnly", BindingFlags.NonPublic | BindingFlags.Instance);
+
+            MethodInfo transcribeOnlyGetMethod = transcribeOnlyProperty.GetGetMethod(nonPublic: true);
+
+            s_getTranscribeOnlyDelegate = (Func<PSHostUserInterface, bool>)Delegate.CreateDelegate(
+                typeof(Func<PSHostUserInterface, bool>), transcribeOnlyGetMethod);
+
+            MethodInfo transcribeOnlySetMethod = transcribeOnlyProperty.GetSetMethod(nonPublic: true);
+
+            s_setTranscribeOnlyDelegate = (Action<PSHostUserInterface, bool>)Delegate.CreateDelegate(
+                typeof(Action<PSHostUserInterface, bool>), transcribeOnlySetMethod);
+
+            s_executionContextProperty = typeof(System.Management.Automation.Runspaces.Runspace)
+                .GetProperty("ExecutionContext", BindingFlags.NonPublic | BindingFlags.Instance);
+
+            s_internalHostProperty = s_executionContextProperty.PropertyType
+                .GetProperty("InternalHost", BindingFlags.NonPublic | BindingFlags.Instance);
+#endif
         }
 
         public PsesInternalHost(
@@ -476,19 +522,7 @@ public void InvokeDelegate(string representation, ExecutionOptions executionOpti
         public IReadOnlyList<TResult> InvokePSCommand<TResult>(PSCommand psCommand, PowerShellExecutionOptions executionOptions, CancellationToken cancellationToken)
         {
             SynchronousPowerShellTask<TResult> task = new(_logger, this, psCommand, executionOptions, cancellationToken);
-            try
-            {
-                return task.ExecuteAndGetResult(cancellationToken);
-            }
-            finally
-            {
-                // At the end of each PowerShell command we need to reset PowerShell 5.1's
-                // `TranscribeOnly` property to avoid a bug where output disappears.
-                if (UI is EditorServicesConsolePSHostUserInterface ui)
-                {
-                    ui.DisableTranscribeOnly();
-                }
-            }
+            return task.ExecuteAndGetResult(cancellationToken);
         }
 
         public void InvokePSCommand(PSCommand psCommand, PowerShellExecutionOptions executionOptions, CancellationToken cancellationToken) => InvokePSCommand<PSObject>(psCommand, executionOptions, cancellationToken);
@@ -507,6 +541,28 @@ public void InvokePSDelegate(string representation, ExecutionOptions executionOp
 
         internal void AddToHistory(string historyEntry) => _readLineProvider.ReadLine.AddToHistory(historyEntry);
 
+        // This works around a bug in PowerShell 5.1 (that was later fixed) where a running
+        // transcription could cause output to disappear since the `TranscribeOnly` property was
+        // accidentally not reset to false.
+#pragma warning disable CA1822 // Warning to make it static when it's empty for CoreCLR.
+        internal void DisableTranscribeOnly()
+#pragma warning restore CA1822
+        {
+#if !CoreCLR
+            // To fix the TranscribeOnly bug, we have to get the internal UI, which involves a lot
+            // of reflection since we can't always just use PowerShell to execute `$Host.UI`.
+            s_internalPSHostUserInterface ??=
+                (s_internalHostProperty.GetValue(
+                    s_executionContextProperty.GetValue(CurrentPowerShell.Runspace))
+                    as PSHost).UI;
+
+            if (s_getTranscribeOnlyDelegate(s_internalPSHostUserInterface))
+            {
+                s_setTranscribeOnlyDelegate(s_internalPSHostUserInterface, false);
+            }
+#endif
+        }
+
         internal Task LoadHostProfilesAsync(CancellationToken cancellationToken)
         {
             // NOTE: This is a special task run on startup!

test/PowerShellEditorServices.Test.E2E/PowerShellEditorServices.Test.E2E.csproj

@@ -6,6 +6,10 @@
     <IsPackable>false</IsPackable>
   </PropertyGroup>
 
+  <PropertyGroup Condition=" '$(TargetFramework)' != 'net462' ">
+    <DefineConstants>$(DefineConstants);CoreCLR</DefineConstants>
+  </PropertyGroup>
+
   <ItemGroup>
     <PackageReference Include="Microsoft.Extensions.Logging" Version="7.0.0" />
     <PackageReference Include="Microsoft.NET.Test.Sdk" Version="17.6.0" />

test/PowerShellEditorServices.Test/PowerShellEditorServices.Test.csproj

@@ -7,6 +7,10 @@
     <TargetPlatform>x64</TargetPlatform>
   </PropertyGroup>
 
+  <PropertyGroup Condition=" '$(TargetFramework)' != 'net462' ">
+    <DefineConstants>$(DefineConstants);CoreCLR</DefineConstants>
+  </PropertyGroup>
+
   <ItemGroup>
     <ProjectReference Include="..\..\src\PowerShellEditorServices\PowerShellEditorServices.csproj" />
     <ProjectReference Include="..\PowerShellEditorServices.Test.Shared\PowerShellEditorServices.Test.Shared.csproj" />
@@ -27,10 +31,6 @@
     <PackageReference Include="Microsoft.PowerShell.5.ReferenceAssemblies" Version="1.1.0" />
   </ItemGroup>
 
-  <PropertyGroup Condition=" '$(TargetFramework)' != 'net462' ">
-    <DefineConstants>$(DefineConstants);CoreCLR</DefineConstants>
-  </PropertyGroup>
-
   <ItemGroup>
     <PackageReference Include="Microsoft.NET.Test.Sdk" Version="17.6.0" />
     <PackageReference Include="xunit" Version="2.4.2" />