From dcf7422d84231d19556b25fe7d64b38a30d4a9bd Mon Sep 17 00:00:00 2001
From: Andrew Schwartzmeyer <andrew@schwartzmeyer.com>
Date: Fri, 18 Feb 2022 17:38:22 -0800
Subject: [PATCH] Reset progress messages at end of REPL

This seems to work. We needed to hold onto processed progress records,
and then at the end of the REPL mark each as completed and re-write it
(which uses the underlying host and effectively clears it).
---
 ...ditorServicesConsolePSHostUserInterface.cs | 37 ++++++++++++++++++-
 .../PowerShell/Host/PsesInternalHost.cs       |  9 +++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/src/PowerShellEditorServices/Services/PowerShell/Host/EditorServicesConsolePSHostUserInterface.cs b/src/PowerShellEditorServices/Services/PowerShell/Host/EditorServicesConsolePSHostUserInterface.cs
index 193f448a5..fbc02a70f 100644
--- a/src/PowerShellEditorServices/Services/PowerShell/Host/EditorServicesConsolePSHostUserInterface.cs
+++ b/src/PowerShellEditorServices/Services/PowerShell/Host/EditorServicesConsolePSHostUserInterface.cs
@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
 using System.Management.Automation;
@@ -22,6 +23,12 @@ internal class EditorServicesConsolePSHostUserInterface : PSHostUserInterface
 
         private readonly PSHostUserInterface _consoleHostUI;
 
+        /// <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();
+
         public EditorServicesConsolePSHostUserInterface(
             ILoggerFactory loggerFactory,
             IReadLineProvider readLineProvider,
@@ -103,7 +110,35 @@ public override PSCredential PromptForCredential(string caption, string message,
 
         public override void WriteLine(string value) => _underlyingHostUI.WriteLine(value);
 
-        public override void WriteProgress(long sourceId, ProgressRecord record) => _underlyingHostUI.WriteProgress(sourceId, record);
+        public override void WriteProgress(long sourceId, ProgressRecord record)
+        {
+            if (record.RecordType == ProgressRecordType.Completed)
+            {
+                _ = _currentProgressRecords.TryRemove((sourceId, record.ActivityId), out _);
+            }
+            else
+            {
+                _ = _currentProgressRecords.TryAdd((sourceId, record.ActivityId), null);
+            }
+            _underlyingHostUI.WriteProgress(sourceId, record);
+        }
+
+        public void ResetProgress()
+        {
+            // Mark all processed progress records as completed.
+            foreach ((long sourceId, int activityId) in _currentProgressRecords.Keys)
+            {
+                // NOTE: This initializer checks that string is not null nor empty, so it must have
+                // some text in it.
+                ProgressRecord record = new(activityId, "0", "0")
+                {
+                    RecordType = ProgressRecordType.Completed
+                };
+                _underlyingHostUI.WriteProgress(sourceId, record);
+                _currentProgressRecords.Clear();
+            }
+            // TODO: Maybe send the OSC sequence to turn off progress indicator.
+        }
 
         public override void WriteVerboseLine(string message) => _underlyingHostUI.WriteVerboseLine(message);
 
diff --git a/src/PowerShellEditorServices/Services/PowerShell/Host/PsesInternalHost.cs b/src/PowerShellEditorServices/Services/PowerShell/Host/PsesInternalHost.cs
index 6f0ddb663..da263e698 100644
--- a/src/PowerShellEditorServices/Services/PowerShell/Host/PsesInternalHost.cs
+++ b/src/PowerShellEditorServices/Services/PowerShell/Host/PsesInternalHost.cs
@@ -662,6 +662,15 @@ private void DoOneRepl(CancellationToken cancellationToken)
                 UI.WriteErrorLine($"An error occurred while running the REPL loop:{Environment.NewLine}{e}");
                 _logger.LogError(e, "An error occurred while running the REPL loop");
             }
+            finally
+            {
+                // At the end of each REPL we need to complete all progress records so that the
+                // progress indicator is cleared.
+                if (UI is EditorServicesConsolePSHostUserInterface ui)
+                {
+                    ui.ResetProgress();
+                }
+            }
         }
 
         private string GetPrompt(CancellationToken cancellationToken)