Color contamination fix

src/AppInstallerCLICore/ChannelStreams.cpp

@@ -25,6 +25,7 @@ namespace AppInstaller::CLI::Execution
     {
         if (m_enabled && m_VTEnabled)
         {
+            m_VTUpdated = true;
             m_out << sequence;
         }
         return *this;
@@ -34,13 +35,26 @@ namespace AppInstaller::CLI::Execution
     {
         if (m_enabled && m_VTEnabled)
         {
+            m_VTUpdated = true;
             m_out << sequence;
         }
         return *this;
     }
 
-    OutputStream::OutputStream(std::ostream& out, bool enabled, bool VTEnabled) :
-        m_out(out, enabled, VTEnabled) {}
+    void BaseStream::Close()
+    {
+        m_enabled = false;
+        if (m_VTUpdated)
+        {
+            Write(TextFormat::Default, true);
+        }
+    }
+
+    OutputStream::OutputStream(BaseStream& out, bool enabled, bool VTEnabled) :
+        m_out(out),
+        m_enabled(enabled),
+        m_VTEnabled(VTEnabled)
+    {}
 
     void OutputStream::AddFormat(const Sequence& sequence)
     {
@@ -61,27 +75,40 @@ namespace AppInstaller::CLI::Execution
 
     OutputStream& OutputStream::operator<<(std::ostream& (__cdecl* f)(std::ostream&))
     {
-        m_out << f;
+        if (m_enabled)
+        {
+            m_out << f;
+        }
+
         return *this;
     }
 
     OutputStream& OutputStream::operator<<(const Sequence& sequence)
     {
-        m_out << sequence;
-        // An incoming sequence will be valid for 1 "standard" output after this one.
-        // We set this to 2 to make that happen, because when it is 1, we will output
-        // the format for the current OutputStream.
-        m_applyFormatAtOne = 2;
+        if (m_enabled && m_VTEnabled)
+        {
+            m_out << sequence;
+
+            // An incoming sequence will be valid for 1 "standard" output after this one.
+            // We set this to 2 to make that happen, because when it is 1, we will output
+            // the format for the current OutputStream.
+            m_applyFormatAtOne = 2;
+        }
+
         return *this;
     }
 
     OutputStream& OutputStream::operator<<(const ConstructedSequence& sequence)
     {
-        m_out << sequence;
-        // An incoming sequence will be valid for 1 "standard" output after this one.
-        // We set this to 2 to make that happen, because when it is 1, we will output
-        // the format for the current OutputStream.
-        m_applyFormatAtOne = 2;
+        if (m_enabled && m_VTEnabled)
+        {
+            m_out << sequence;
+            // An incoming sequence will be valid for 1 "standard" output after this one.
+            // We set this to 2 to make that happen, because when it is 1, we will output
+            // the format for the current OutputStream.
+            m_applyFormatAtOne = 2;
+        }
+
         return *this;
     }
 
@@ -91,13 +118,4 @@ namespace AppInstaller::CLI::Execution
         m_out << path.u8string();
         return *this;
     }
-
-    NoVTStream::NoVTStream(std::ostream& out, bool enabled) :
-        m_out(out, enabled, false) {}
-
-    NoVTStream& NoVTStream::operator<<(std::ostream& (__cdecl* f)(std::ostream&))
-    {
-        m_out << f;
-        return *this;
-    }
-}
+}

src/AppInstallerCLICore/ChannelStreams.h

@@ -1,7 +1,6 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 #pragma once
-#include "ExecutionProgress.h"
 #include "Resources.h"
 #include "VTSupport.h"
 #include <winget/LocIndependent.h>
@@ -50,27 +49,36 @@ namespace AppInstaller::CLI::Execution
         template <typename T>
         BaseStream& operator<<(const T& t)
         {
-            if (m_enabled)
-            {
-                m_out << t;
-            }
+            Write<T>(t, false);
             return *this;
         }
 
         BaseStream& operator<<(std::ostream& (__cdecl* f)(std::ostream&));
         BaseStream& operator<<(const VirtualTerminal::Sequence& sequence);
         BaseStream& operator<<(const VirtualTerminal::ConstructedSequence& sequence);
 
+        void Close();
+
     private:
+        template <typename T>
+        void Write(const T& t, bool bypass)
+        {
+            if (bypass || m_enabled)
+            {
+                m_out << t;
+            }
+        };
+
         std::ostream& m_out;
-        bool m_enabled;
+        std::atomic_bool m_enabled;
+        std::atomic_bool m_VTUpdated;
         bool m_VTEnabled;
     };
 
     // Holds output formatting information.
     struct OutputStream
     {
-        OutputStream(std::ostream& out, bool enabled, bool VTEnabled);
+        OutputStream(BaseStream& out, bool enabled, bool VTEnabled);
 
         // Adds a format to the current value.
         void AddFormat(const VirtualTerminal::Sequence& sequence);
@@ -92,7 +100,10 @@ namespace AppInstaller::CLI::Execution
             //      informs the output that there is no localized version to use.
             // TODO: Convert the rest of the code base and uncomment to enforce localization.
             //static_assert(details::IsApprovedForOutput<std::decay_t<T>>::value, "This type may not be localized, see comment for more information");
-            ApplyFormat();
+            if (m_VTEnabled)
+            {
+                ApplyFormat();
+            }
             m_out << t;
             return *this;
         }
@@ -106,28 +117,10 @@ namespace AppInstaller::CLI::Execution
         // Applies the format for the stream.
         void ApplyFormat();
 
-        BaseStream m_out;
+        BaseStream& m_out;
+        bool m_enabled;
+        bool m_VTEnabled;
         size_t m_applyFormatAtOne = 1;
         VirtualTerminal::ConstructedSequence m_format;
     };
-
-    // Does not allow VT at all.
-    struct NoVTStream
-    {
-        NoVTStream(std::ostream& out, bool enabled);
-
-        template <typename T>
-        NoVTStream& operator<<(const T& t)
-        {
-            m_out << t;
-            return *this;
-        }
-
-        NoVTStream& operator<<(std::ostream& (__cdecl* f)(std::ostream&));
-        NoVTStream& operator<<(const VirtualTerminal::Sequence& sequence) = delete;
-        NoVTStream& operator<<(const VirtualTerminal::ConstructedSequence& sequence) = delete;
-
-    private:
-        BaseStream m_out;
-    };
 }

src/AppInstallerCLICore/ExecutionContext.cpp

@@ -169,6 +169,7 @@ namespace AppInstaller::CLI::Execution
             // Unless we want to spin a separate thread for all work, we have to just exit here.
             if (m_CtrlSignalCount >= 2)
             {
+                Reporter.CloseOutputStream();
                 Logging::Telemetry().LogCommandTermination(hr, file, line);
                 std::exit(hr);
             }

src/AppInstallerCLICore/ExecutionProgress.cpp

@@ -41,7 +41,7 @@ namespace AppInstaller::CLI::Execution
             return s_bytesFormatData[ARRAYSIZE(s_bytesFormatData) - 1];
         }
 
-        void OutputBytes(std::ostream& out, uint64_t byteCount)
+        void OutputBytes(BaseStream& out, uint64_t byteCount)
         {
             const BytesFormatData& bfd = GetFormatForSize(byteCount);
 
@@ -76,7 +76,7 @@ namespace AppInstaller::CLI::Execution
             out << ' ' << bfd.Name;
         }
 
-        void SetColor(std::ostream& out, const TextFormat::Color& color, bool enabled)
+        void SetColor(BaseStream& out, const TextFormat::Color& color, bool enabled)
         {
             if (enabled)
             {
@@ -95,7 +95,7 @@ namespace AppInstaller::CLI::Execution
             }
         }
 
-        void SetRainbowColor(std::ostream& out, size_t i, size_t max, bool enabled)
+        void SetRainbowColor(BaseStream& out, size_t i, size_t max, bool enabled)
         {
             TextFormat::Color rainbow[] =
             {

src/AppInstallerCLICore/ExecutionProgress.h

@@ -4,6 +4,7 @@
 #include "VTSupport.h"
 #include <AppInstallerProgress.h>
 #include <winget/UserSettings.h>
+#include <ChannelStreams.h>
 
 #include <wil/resource.h>
 
@@ -21,13 +22,13 @@ namespace AppInstaller::CLI::Execution
         // Shared functionality for progress visualizers.
         struct ProgressVisualizerBase
         {
-            ProgressVisualizerBase(std::ostream& stream, bool enableVT) :
+            ProgressVisualizerBase(BaseStream& stream, bool enableVT) :
                 m_out(stream), m_enableVT(enableVT) {}
 
             void SetStyle(AppInstaller::Settings::VisualStyle style) { m_style = style; }
 
         protected:
-            std::ostream& m_out;
+            BaseStream& m_out;
             Settings::VisualStyle m_style = AppInstaller::Settings::VisualStyle::Accent;
 
             bool UseVT() const { return m_enableVT && m_style != AppInstaller::Settings::VisualStyle::NoVT; }
@@ -43,7 +44,7 @@ namespace AppInstaller::CLI::Execution
     // Displays an indefinite spinner.
     struct IndefiniteSpinner : public details::ProgressVisualizerBase
     {
-        IndefiniteSpinner(std::ostream& stream, bool enableVT) :
+        IndefiniteSpinner(BaseStream& stream, bool enableVT) :
             details::ProgressVisualizerBase(stream, enableVT) {}
 
         void ShowSpinner();
@@ -62,7 +63,7 @@ namespace AppInstaller::CLI::Execution
     class ProgressBar : public details::ProgressVisualizerBase
     {
     public:
-        ProgressBar(std::ostream& stream, bool enableVT) :
+        ProgressBar(BaseStream& stream, bool enableVT) :
             details::ProgressVisualizerBase(stream, enableVT) {}
 
         void ShowProgress(uint64_t current, uint64_t maximum, ProgressType type);

src/AppInstallerCLICore/ExecutionReporter.cpp

@@ -19,19 +19,23 @@ namespace AppInstaller::CLI::Execution
     const Sequence& PromptEmphasis = TextFormat::Foreground::Bright;
 
     Reporter::Reporter(std::ostream& outStream, std::istream& inStream) :
+        Reporter(std::make_shared<BaseStream>(outStream, true, IsVTEnabled()), inStream)
+    {
+        SetProgressSink(this);
+    }
+
+    Reporter::Reporter(std::shared_ptr<BaseStream> outStream, std::istream& inStream) :
         m_out(outStream),
         m_in(inStream),
-        m_progressBar(std::in_place, outStream, IsVTEnabled()),
-        m_spinner(std::in_place, outStream, IsVTEnabled())
+        m_progressBar(std::in_place, *m_out, IsVTEnabled()),
+        m_spinner(std::in_place, *m_out, IsVTEnabled())
     {
         SetProgressSink(this);
     }
 
     Reporter::~Reporter()
     {
-        // The goal of this is to return output to its previous state.
-        // For now, we assume this means "default".
-        GetBasicOutputStream() << TextFormat::Default;
+        this->CloseOutputStream();
     }
 
     Reporter::Reporter(const Reporter& other, clone_t) :
@@ -70,7 +74,7 @@ namespace AppInstaller::CLI::Execution
 
     OutputStream Reporter::GetBasicOutputStream()
     {
-        return { m_out, m_channel == Channel::Output, IsVTEnabled() };
+        return {*m_out, m_channel == Channel::Output, IsVTEnabled() };
     }
 
     void Reporter::SetChannel(Channel channel)
@@ -213,4 +217,9 @@ namespace AppInstaller::CLI::Execution
     {
         return m_isVTEnabled && ConsoleModeRestore::Instance().IsVTEnabled();
     }
-}
+
+    void Reporter::CloseOutputStream()
+    {
+        m_out->Close();
+    }
+}

src/AppInstallerCLICore/ExecutionReporter.h

@@ -81,7 +81,7 @@ namespace AppInstaller::CLI::Execution
         OutputStream Error() { return GetOutputStream(Level::Error); }
 
         // Get a stream for outputting completion words.
-        NoVTStream Completion() { return NoVTStream(m_out, m_channel == Channel::Completion); }
+        OutputStream Completion() { return OutputStream(*m_out, m_channel == Channel::Completion, false); }
 
         // Gets a stream for output of the given level.
         OutputStream GetOutputStream(Level level);
@@ -131,20 +131,23 @@ namespace AppInstaller::CLI::Execution
         // Cancels the in progress task.
         void CancelInProgressTask(bool force);
 
+        void CloseOutputStream();
+
         void SetProgressSink(IProgressSink* sink)
         {
             m_progressSink = sink;
         }
 
     private:
+        Reporter(std::shared_ptr<BaseStream> outStream, std::istream& inStream);
         // Gets whether VT is enabled for this reporter.
         bool IsVTEnabled() const;
 
         // Gets a stream for output for internal use.
         OutputStream GetBasicOutputStream();
 
         Channel m_channel = Channel::Output;
-        std::ostream& m_out;
+        std::shared_ptr<BaseStream> m_out;
         std::istream& m_in;
         bool m_isVTEnabled = true;
         std::optional<AppInstaller::Settings::VisualStyle> m_style;

src/AppInstallerCLICore/Workflows/CompletionFlow.cpp

@@ -11,7 +11,7 @@ namespace AppInstaller::CLI::Workflow
     namespace
     {
         // Outputs the completion string, wrapping it in quotes if needed.
-        void OutputCompletionString(Execution::NoVTStream& stream, std::string_view value)
+        void OutputCompletionString(Execution::OutputStream& stream, std::string_view value)
         {
             if (value.find_first_of(' ') != std::string_view::npos)
             {