Use NotifyEndApplication to re-enable VT mode (PowerShell#16612)

Use NotifyEndApplication to re-enable VT mode instead of doing it in InputLoop.Run, so that we run this code only after invoking native commands

Author: daxian-dbw

src/Microsoft.PowerShell.ConsoleHost/host/msh/ConsoleHost.cs

@@ -1071,18 +1071,16 @@ public override void NotifyBeginApplication()
         {
             lock (hostGlobalLock)
             {
-                ++_beginApplicationNotifyCount;
-                if (_beginApplicationNotifyCount == 1)
+                if (++_beginApplicationNotifyCount == 1)
                 {
-                    // save the window title when first notified.
-
+                    // Save the window title when first notified.
                     _savedWindowTitle = ui.RawUI.WindowTitle;
 #if !UNIX
                     if (_initialConsoleMode != ConsoleControl.ConsoleModes.Unknown)
                     {
-                        var activeScreenBufferHandle = ConsoleControl.GetActiveScreenBufferHandle();
-                        _savedConsoleMode = ConsoleControl.GetMode(activeScreenBufferHandle);
-                        ConsoleControl.SetMode(activeScreenBufferHandle, _initialConsoleMode);
+                        var outputHandle = ConsoleControl.GetActiveScreenBufferHandle();
+                        _savedConsoleMode = ConsoleControl.GetMode(outputHandle);
+                        ConsoleControl.SetMode(outputHandle, _initialConsoleMode);
                     }
 #endif
                 }
@@ -1097,17 +1095,26 @@ public override void NotifyEndApplication()
         {
             lock (hostGlobalLock)
             {
-                Dbg.Assert(_beginApplicationNotifyCount > 0, "Not running an executable - NotifyBeginApplication was not called!");
-                --_beginApplicationNotifyCount;
-                if (_beginApplicationNotifyCount == 0)
+                if (--_beginApplicationNotifyCount == 0)
                 {
-                    // restore the window title when the last application started has ended.
-
+                    // Restore the window title when the last application started has ended.
                     ui.RawUI.WindowTitle = _savedWindowTitle;
 #if !UNIX
                     if (_savedConsoleMode != ConsoleControl.ConsoleModes.Unknown)
                     {
                         ConsoleControl.SetMode(ConsoleControl.GetActiveScreenBufferHandle(), _savedConsoleMode);
+                        if (_savedConsoleMode.HasFlag(ConsoleControl.ConsoleModes.VirtualTerminal))
+                        {
+                            // If the console output mode we just set already has 'VirtualTerminal' turned on,
+                            // we don't need to try turn on the VT mode separately.
+                            return;
+                        }
+                    }
+
+                    if (ui.SupportsVirtualTerminal)
+                    {
+                        // Re-enable VT mode if it was previously enabled, as a native command may have turned it off.
+                        ui.TryTurnOnVirtualTerminal();
                     }
 #endif
                 }
@@ -2452,14 +2459,6 @@ internal void Run(bool inputLoopIsNested)
 
                 while (!_parent.ShouldEndSession && !_shouldExit)
                 {
-#if !UNIX
-                    if (ui.SupportsVirtualTerminal)
-                    {
-                        // need to re-enable VT mode if it was previously enabled as native commands may have turned it off
-                        ui.TryTurnOnVtMode();
-                    }
-#endif
-
                     try
                     {
                         _parent._isRunningPromptLoop = true;

src/Microsoft.PowerShell.ConsoleHost/host/msh/ConsoleHostUserInterface.cs

@@ -87,28 +87,34 @@ internal ConsoleHostUserInterface(ConsoleHost parent)
 
             if (SupportsVirtualTerminal)
             {
-                SupportsVirtualTerminal = TryTurnOnVtMode();
+                SupportsVirtualTerminal = TryTurnOnVirtualTerminal();
             }
         }
 
-        internal bool TryTurnOnVtMode()
+        internal bool TryTurnOnVirtualTerminal()
         {
 #if UNIX
             return true;
 #else
             try
             {
                 // Turn on virtual terminal if possible.
-
                 // This might throw - not sure how exactly (no console), but if it does, we shouldn't fail to start.
-                var handle = ConsoleControl.GetActiveScreenBufferHandle();
-                var m = ConsoleControl.GetMode(handle);
-                if (ConsoleControl.NativeMethods.SetConsoleMode(handle.DangerousGetHandle(), (uint)(m | ConsoleControl.ConsoleModes.VirtualTerminal)))
+                var outputHandle = ConsoleControl.GetActiveScreenBufferHandle();
+                var outputMode = ConsoleControl.GetMode(outputHandle);
+
+                if (outputMode.HasFlag(ConsoleControl.ConsoleModes.VirtualTerminal))
+                {
+                    return true;
+                }
+
+                outputMode |= ConsoleControl.ConsoleModes.VirtualTerminal;
+                if (ConsoleControl.NativeMethods.SetConsoleMode(outputHandle.DangerousGetHandle(), (uint)outputMode))
                 {
                     // We only know if vt100 is supported if the previous call actually set the new flag, older
                     // systems ignore the setting.
-                    m = ConsoleControl.GetMode(handle);
-                    return (m & ConsoleControl.ConsoleModes.VirtualTerminal) != 0;
+                    outputMode = ConsoleControl.GetMode(outputHandle);
+                    return outputMode.HasFlag(ConsoleControl.ConsoleModes.VirtualTerminal);
                 }
             }
             catch

src/System.Management.Automation/engine/NativeCommandProcessor.cs

@@ -545,15 +545,14 @@ private void InitNativeProcess()
             Exception exceptionToRethrow = null;
             try
             {
-                // If this process is being run standalone, tell the host, which might want
-                // to save off the window title or other such state as might be tweaked by
-                // the native process
+                // Before start the executable, tell the host, which might want to save off the
+                // window title or other such state as might be tweaked by the native process.
+                Command.Context.EngineHostInterface.NotifyBeginApplication();
+                _hasNotifiedBeginApplication = true;
+
                 if (_runStandAlone)
                 {
-                    this.Command.Context.EngineHostInterface.NotifyBeginApplication();
-                    _hasNotifiedBeginApplication = true;
-
-                    // Also, store the Raw UI coordinates so that we can scrape the screen after
+                    // Store the Raw UI coordinates so that we can scrape the screen after
                     // if we are transcribing.
                     if (_isTranscribing && (s_supportScreenScrape == true))
                     {
@@ -1166,7 +1165,7 @@ private void CleanUp()
             // We need to call 'NotifyEndApplication' as appropriate during cleanup
             if (_hasNotifiedBeginApplication)
             {
-                this.Command.Context.EngineHostInterface.NotifyEndApplication();
+                Command.Context.EngineHostInterface.NotifyEndApplication();
             }
 
             try