Fix the TranscribeOnly bug (take two)

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

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/Services/PowerShell/Execution/SynchronousPowerShellTask.cs

@@ -5,7 +5,9 @@
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
 using System.Management.Automation;
+using System.Management.Automation.Host;
 using System.Management.Automation.Remoting;
+using System.Reflection;
 using System.Threading;
 using Microsoft.Extensions.Logging;
 using Microsoft.PowerShell.EditorServices.Services.PowerShell.Context;
@@ -27,6 +29,55 @@ internal class SynchronousPowerShellTask<TResult> : SynchronousTask<IReadOnlyLis
     {
         private static readonly PowerShellExecutionOptions s_defaultPowerShellExecutionOptions = new();
 
+#if !CoreCLR
+        /// <summary>
+        /// To workaround a bug where the `TranscribeOnly` field of the PSHostUserInterface can
+        /// accidentally remain true, we have to use a PowerShell pipeline to get the internal
+        /// instance (saved here) and reflection to get delegates for the field's getter and setter
+        /// methods, and then reset it to false (see <see cref="ExecuteNormally"/>). Note that it
+        /// must be the internal instance, not our own UI instance.
+        /// See https://github.com/PowerShell/PowerShell/pull/3436
+        /// </summary>
+        [ThreadStatic] // Because we can re-use it, but only for each PowerShell.
+        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;
+
+        private static readonly PropertyInfo s_internalHostUIProperty;
+
+        static SynchronousPowerShellTask()
+        {
+            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(SMA.Runspaces.Runspace)
+                .GetProperty("ExecutionContext", BindingFlags.NonPublic | BindingFlags.Instance);
+
+            s_internalHostProperty = s_executionContextProperty.PropertyType
+                .GetProperty("InternalHost", BindingFlags.NonPublic | BindingFlags.Instance);
+
+            // It's public but we want the override and reflection confuses me.
+            s_internalHostUIProperty = s_internalHostProperty.PropertyType
+                .GetProperty("UI", BindingFlags.Public | BindingFlags.Instance);
+        }
+#endif
+
         private readonly ILogger _logger;
 
         private readonly PsesInternalHost _psesHost;
@@ -105,6 +156,21 @@ private IReadOnlyList<TResult> ExecuteNormally(CancellationToken cancellationTok
             if (PowerShellExecutionOptions.WriteOutputToHost)
             {
                 _psCommand.AddOutputCommand();
+#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`
+                // (we tried). With that internal UI we can reset its `TranscribeOnly` flag and so
+                // avoid the disappearing output that happens when a transcription is running.
+                if (!_pwsh.Runspace.RunspaceIsRemote)
+                {
+                    s_internalPSHostUserInterface ??=
+                        s_internalHostUIProperty.GetValue(
+                            s_internalHostProperty.GetValue(
+                                s_executionContextProperty.GetValue(_pwsh.Runspace)))
+                        as PSHostUserInterface;
+                    DisableTranscribeOnly();
+                }
+#endif
             }
 
             cancellationToken.Register(CancelNormalExecution);
@@ -148,7 +214,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 +255,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 +311,7 @@ private IReadOnlyList<TResult> ExecuteInDebugger(CancellationToken cancellationT
                 if (e is PSRemotingTransportException)
                 {
                     _ = System.Threading.Tasks.Task.Run(
-                        () => _psesHost.UnwindCallStack(),
+                        _psesHost.UnwindCallStack,
                         CancellationToken.None)
                         .HandleErrorsAsync(_logger);
 
@@ -396,5 +460,18 @@ private void CancelDebugExecution()
 
             _pwsh.Runspace.Debugger.StopProcessCommand();
         }
+
+#if !CoreCLR
+        // 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.
+        private static void DisableTranscribeOnly()
+        {
+            if (s_getTranscribeOnlyDelegate(s_internalPSHostUserInterface))
+            {
+                s_setTranscribeOnlyDelegate(s_internalPSHostUserInterface, false);
+            }
+        }
+#endif
     }
 }

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

@@ -476,19 +476,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);