From bc998901e41e7da976079b96a5701d6107ef61f8 Mon Sep 17 00:00:00 2001
From: Andy Jordan <andy.jordan@microsoft.com>
Date: Tue, 2 Aug 2022 13:27:01 -0700
Subject: [PATCH] Fix bug where error in `prompt` function crashed REPL loop

Double-whammy fix by both setting `ThrowOnError` to false, and catching
a more generic PowerShell `RuntimeException`. Plus a regression test!
---
 .../PowerShell/Host/PsesInternalHost.cs       | 12 ++++++----
 .../Session/PsesInternalHostTests.cs          | 24 ++++++++++++++++---
 2 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/src/PowerShellEditorServices/Services/PowerShell/Host/PsesInternalHost.cs b/src/PowerShellEditorServices/Services/PowerShell/Host/PsesInternalHost.cs
index d259a8ac9..e1ee99345 100644
--- a/src/PowerShellEditorServices/Services/PowerShell/Host/PsesInternalHost.cs
+++ b/src/PowerShellEditorServices/Services/PowerShell/Host/PsesInternalHost.cs
@@ -31,7 +31,7 @@ namespace Microsoft.PowerShell.EditorServices.Services.PowerShell.Host
 
     internal class PsesInternalHost : PSHost, IHostSupportsInteractiveSession, IRunspaceContext, IInternalPowerShellExecutionService
     {
-        private const string DefaultPrompt = "> ";
+        internal const string DefaultPrompt = "> ";
 
         private static readonly PropertyInfo s_scriptDebuggerTriggerObjectProperty;
 
@@ -865,7 +865,7 @@ private void DoOneRepl(CancellationToken cancellationToken)
             }
         }
 
-        private string GetPrompt(CancellationToken cancellationToken)
+        internal string GetPrompt(CancellationToken cancellationToken)
         {
             Runspace.ThrowCancelledIfUnusable();
             string prompt = DefaultPrompt;
@@ -873,13 +873,17 @@ private string GetPrompt(CancellationToken cancellationToken)
             {
                 // TODO: Should we cache PSCommands like this as static members?
                 PSCommand command = new PSCommand().AddCommand("prompt");
-                IReadOnlyList<string> results = InvokePSCommand<string>(command, executionOptions: null, cancellationToken);
+                IReadOnlyList<string> results = InvokePSCommand<string>(
+                    command,
+                    executionOptions: new PowerShellExecutionOptions { ThrowOnError = false },
+                    cancellationToken);
+
                 if (results?.Count > 0)
                 {
                     prompt = results[0];
                 }
             }
-            catch (CommandNotFoundException) { } // Use default prompt
+            catch (RuntimeException) { } // Use default prompt
 
             if (CurrentRunspace.RunspaceOrigin != RunspaceOrigin.Local)
             {
diff --git a/test/PowerShellEditorServices.Test/Session/PsesInternalHostTests.cs b/test/PowerShellEditorServices.Test/Session/PsesInternalHostTests.cs
index 45f14e97c..41d5cbb82 100644
--- a/test/PowerShellEditorServices.Test/Session/PsesInternalHostTests.cs
+++ b/test/PowerShellEditorServices.Test/Session/PsesInternalHostTests.cs
@@ -162,12 +162,30 @@ await psesHost.ExecuteDelegateAsync(
                 CancellationToken.None).ConfigureAwait(true);
         }
 
+        // NOTE: Tests where we call functions that use PowerShell runspaces are slightly more
+        // complicated than one would expect because we explicitly need the methods to run on the
+        // pipeline thread, otherwise Windows complains about the the thread's apartment state not
+        // matching. Hence we use a delegate where it looks like we could just call the method.
+
+        [Fact]
+        public async Task CanHandleBrokenPrompt()
+        {
+            await psesHost.ExecutePSCommandAsync(
+                new PSCommand().AddScript("function prompt { throw }"),
+                CancellationToken.None).ConfigureAwait(true);
+
+            string prompt = await psesHost.ExecuteDelegateAsync(
+                nameof(psesHost.GetPrompt),
+                executionOptions: null,
+                (_, _) => psesHost.GetPrompt(CancellationToken.None),
+                CancellationToken.None).ConfigureAwait(true);
+
+            Assert.Equal(PsesInternalHost.DefaultPrompt, prompt);
+        }
+
         [Fact]
         public async Task CanLoadPSReadLine()
         {
-            // NOTE: This is slightly more complicated than one would expect because we explicitly
-            // need it to run on the pipeline thread otherwise Windows complains about the the
-            // thread's apartment state not matching.
             Assert.True(await psesHost.ExecuteDelegateAsync(
                 nameof(psesHost.TryLoadPSReadLine),
                 executionOptions: null,