From d6efcabcc14c044e285a12eb9b4a6ef11fa2a657 Mon Sep 17 00:00:00 2001
From: Dongbo Wang <dongbow@microsoft.com>
Date: Wed, 29 Mar 2023 11:27:55 -0700
Subject: [PATCH] Set the current location in `PredictionClient` when it's
 supported (#3639)

---
 PSReadLine/Prediction.cs | 40 ++++++++++++++++++++++++++++++++++++++--
 PSReadLine/ReadLine.cs   | 10 +++++-----
 2 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/PSReadLine/Prediction.cs b/PSReadLine/Prediction.cs
index dfa1f1c..fc790f9 100644
--- a/PSReadLine/Prediction.cs
+++ b/PSReadLine/Prediction.cs
@@ -4,8 +4,10 @@ Copyright (c) Microsoft Corporation.  All rights reserved.
 
 using System;
 using System.Collections.Generic;
+using System.Reflection;
 using System.Threading.Tasks;
 using System.Management.Automation;
+using System.Management.Automation.Runspaces;
 using System.Management.Automation.Language;
 using System.Management.Automation.Subsystem.Prediction;
 using System.Diagnostics.CodeAnalysis;
@@ -16,8 +18,42 @@ namespace Microsoft.PowerShell
 {
     public partial class PSConsoleReadLine
     {
-        private const string PSReadLine = "PSReadLine";
-        private static PredictionClient s_predictionClient = new(PSReadLine, PredictionClientKind.Terminal);
+        private const string DefaultName = "PSReadLine";
+        private static readonly PredictionClient s_predictionClient = new(DefaultName, PredictionClientKind.Terminal);
+        private static PropertyInfo s_pCurrentLocation = null;
+
+        /// <summary>
+        /// Initialize the <see cref="PropertyInfo"/> objects for those public settable properties newly added to
+        /// <see cref="PredictionClient"/>.
+        /// </summary>
+        private static void InitializePropertyInfo()
+        {
+            Version ver = typeof(PSObject).Assembly.GetName().Version;
+            if (ver.Major < 7 || ver.Minor < 4)
+            {
+                return;
+            }
+
+            Type pcType = typeof(PredictionClient);
+            // Property added in 7.4
+            s_pCurrentLocation = pcType.GetProperty("CurrentLocation");
+        }
+
+        /// <summary>
+        /// New public settable properties may be added to the <see cref="PredictionClient"/> type as it evolves to
+        /// offer more helpful context information. We dynamically set those properties here to avoid any backward
+        /// compatibility issues.
+        /// </summary>
+        private static void UpdatePredictionClient(Runspace runspace, EngineIntrinsics engineIntrinsics)
+        {
+            // Set the current location if the 'CurrentLocation' property exists.
+            if (s_pCurrentLocation is not null)
+            {
+                // Set the current location if it's a local Runspace. Otherwise, set it to null.
+                object path = runspace.RunspaceIsRemote ? null : engineIntrinsics.SessionState.Path.CurrentLocation;
+                s_pCurrentLocation.SetValue(s_predictionClient, path);
+            }
+        }
 
         // Stub helper methods so prediction can be mocked
         [ExcludeFromCodeCoverage]
diff --git a/PSReadLine/ReadLine.cs b/PSReadLine/ReadLine.cs
index 33b65d0..3b042a4 100644
--- a/PSReadLine/ReadLine.cs
+++ b/PSReadLine/ReadLine.cs
@@ -652,6 +652,7 @@ namespace Microsoft.PowerShell
         {
             _singleton = new PSConsoleReadLine();
             _viRegister = new ViRegister(_singleton);
+            InitializePropertyInfo();
         }
 
         private PSConsoleReadLine()
@@ -682,13 +683,9 @@ namespace Microsoft.PowerShell
                 {
                 }
             }
-            if (hostName == null)
-            {
-                hostName = PSReadLine;
-            }
 
             bool usingLegacyConsole = _console is PlatformWindows.LegacyWin32Console;
-            _options = new PSConsoleReadLineOptions(hostName, usingLegacyConsole);
+            _options = new PSConsoleReadLineOptions(hostName ?? DefaultName, usingLegacyConsole);
             _prediction = new Prediction(this);
             SetDefaultBindings(_options.EditMode);
         }
@@ -698,6 +695,9 @@ namespace Microsoft.PowerShell
             _engineIntrinsics = engineIntrinsics;
             _runspace = runspace;
 
+            // Update the client instance per every call to PSReadLine.
+            UpdatePredictionClient(runspace, engineIntrinsics);
+
             if (!_delayedOneTimeInitCompleted)
             {
                 DelayedOneTimeInitialize();