From 09b90d16cdff86a4e9896c115f97b0d2725aab99 Mon Sep 17 00:00:00 2001
From: David Wilson <daviwil@microsoft.com>
Date: Tue, 14 Mar 2017 12:32:41 -0700
Subject: [PATCH] Fix remote prompt ComputerName display

This change fixes the ComputerName prefix for prompts shown in remote
sessions.  Recent changes caused it to be lost.
---
 .../Console/ConsoleService.cs                 | 21 ++++++++++++++++---
 .../Session/PowerShellContext.cs              | 11 ----------
 2 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/src/PowerShellEditorServices/Console/ConsoleService.cs b/src/PowerShellEditorServices/Console/ConsoleService.cs
index 64b6b2b8f..5efa1a4d3 100644
--- a/src/PowerShellEditorServices/Console/ConsoleService.cs
+++ b/src/PowerShellEditorServices/Console/ConsoleService.cs
@@ -10,7 +10,9 @@
 
 namespace Microsoft.PowerShell.EditorServices.Console
 {
+    using Microsoft.PowerShell.EditorServices.Session;
     using System;
+    using System.Globalization;
     using System.Management.Automation;
     using System.Security;
 
@@ -264,10 +266,23 @@ public async Task<SecureString> ReadSecureLine(CancellationToken cancellationTok
 
         private void WritePromptStringToHost()
         {
+            string promptString = this.powerShellContext.PromptString;
+
+            // Update the stored prompt string if the session is remote
+            if (this.powerShellContext.CurrentRunspace.Location == RunspaceLocation.Remote)
+            {
+                promptString =
+                    string.Format(
+                        CultureInfo.InvariantCulture,
+                        "[{0}]: {1}",
+                        this.powerShellContext.CurrentRunspace.Runspace.ConnectionInfo != null
+                            ? this.powerShellContext.CurrentRunspace.Runspace.ConnectionInfo.ComputerName
+                            : this.powerShellContext.CurrentRunspace.SessionDetails.ComputerName,
+                        promptString);
+            }
+
             // Write the prompt string
-            this.WriteOutput(
-                this.powerShellContext.PromptString,
-                false);
+            this.WriteOutput(promptString, false);
         }
 
         private void WriteDebuggerBanner(DebuggerStopEventArgs eventArgs)
diff --git a/src/PowerShellEditorServices/Session/PowerShellContext.cs b/src/PowerShellEditorServices/Session/PowerShellContext.cs
index dd4e1f76d..d4afb44db 100644
--- a/src/PowerShellEditorServices/Session/PowerShellContext.cs
+++ b/src/PowerShellEditorServices/Session/PowerShellContext.cs
@@ -1367,17 +1367,6 @@ private SessionDetails GetSessionDetailsInRunspace(Runspace runspace)
                         }
                     });
 
-            if (this.CurrentRunspace != null &&
-                this.CurrentRunspace.Location == RunspaceLocation.Remote)
-            {
-                sessionDetails.PromptString =
-                    string.Format(
-                        CultureInfo.InvariantCulture,
-                        "[{0}]: {1}",
-                        runspace.ConnectionInfo.ComputerName,
-                        sessionDetails.PromptString);
-            }
-
             return sessionDetails;
         }