Improve the telemetry and http request

tools/Az.Tools.Predictor/Az.Tools.Predictor.Test/AzPredictorTests.cs

@@ -45,7 +45,8 @@ public AzPredictorTests(ModelFixture modelFixture)
             this._azPredictor = new AzPredictor(this._service, this._telemetryClient, new Settings()
             {
                 SuggestionCount = 1,
-            });
+            },
+            null);
         }
 
         /// <summary>

tools/Az.Tools.Predictor/Az.Tools.Predictor/Az.Tools.Predictor.csproj

@@ -33,6 +33,7 @@
 
   <ItemGroup>
     <PackageReference Include="Microsoft.ApplicationInsights" Version="2.14.0" />
+    <PackageReference Include="Microsoft.PowerShell.SDK" Version="7.1.0-preview.7" />
     <PackageReference Include="System.Management.Automation" Version="7.1.0-preview.7" />
   </ItemGroup>
 

tools/Az.Tools.Predictor/Az.Tools.Predictor/AzContext.cs

@@ -0,0 +1,177 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+
+using System.Management.Automation;
+using System;
+using System.Collections.Generic;
+using System.Collections.ObjectModel;
+using System.Linq;
+using System.Net.NetworkInformation;
+using System.Security.Cryptography;
+using System.Text;
+
+namespace Microsoft.Azure.PowerShell.Tools.AzPredictor
+{
+    using PowerShell = System.Management.Automation.PowerShell;
+
+    /// <summary>
+    /// The class for the current Azure PowerShell context.
+    /// </summary>
+    internal sealed class AzContext : IAzContext
+    {
+        /// <inheritdoc/>
+        public Version AzVersion { get; private set; } = new Version("0.0.0.0");
+
+        /// <inheritdoc/>
+        public string HashedUserId { get; private set; } = string.Empty;
+
+        private string _cachedHashedMacAddress;
+        /// <inheritdoc/>
+        public string HashedMacAddress
+        {
+            get
+            {
+                if (_cachedHashedMacAddress == null)
+                {
+                    _cachedHashedMacAddress = string.Empty;
+
+                    var macAddress = GetMACAddress();
+                    if (!string.IsNullOrWhiteSpace(macAddress))
+                    {
+                        _cachedHashedMacAddress = GenerateSha256HashString(macAddress)?.Replace("-", string.Empty).ToLowerInvariant();
+                    }
+                }
+
+                return _cachedHashedMacAddress;
+            }
+        }
+
+        /// <inheritdoc/>
+        public void UpdateContext()
+        {
+            AzVersion = GetAzVersion();
+            HashedUserId = GenerateSha256HashString(GetUserAccountId());
+        }
+
+        /// <summary>
+        /// Gets the latest version from the loaded Az modules.
+        /// </summary>
+        private Version GetAzVersion()
+        {
+            Version defaultVersion = new Version("0.0.0");
+
+            Version latestAz = defaultVersion;
+
+            try
+            {
+                var outputs = AzContext.ExecuteScript<PSObject>("Get-Module -Name Az -ListAvailable");
+                foreach (PSObject obj in outputs)
+                {
+                    string psVersion = obj.Properties["Version"].Value.ToString();
+                    int pos = psVersion.IndexOf('-');
+                    Version currentAz = (pos == -1) ? new Version(psVersion) : new Version(psVersion.Substring(0, pos));
+                    if (currentAz > latestAz)
+                    {
+                        latestAz = currentAz;
+                    }
+                }
+            }
+            catch (Exception)
+            {
+            }
+
+            return latestAz;
+        }
+
+        /// <summary>
+        /// Gets the user account id if the user logs in, otherwise empty string.
+        /// </summary>
+        private string GetUserAccountId()
+        {
+            try
+            {
+                var output = AzContext.ExecuteScript<string>("(Get-AzContext).Account.Id");
+                return output.FirstOrDefault() ?? string.Empty;
+            }
+            catch (Exception)
+            {
+            }
+
+            return string.Empty;
+        }
+
+        /// <summary>
+        /// Executes the PowerShell cmdlet in the current powershell session.
+        /// </summary>
+        private static List<T> ExecuteScript<T>(string contents)
+        {
+            List<T> output = new List<T>();
+
+            using (PowerShell powershell = PowerShell.Create(RunspaceMode.NewRunspace))
+            {
+                powershell.AddScript(contents);
+                Collection<T> result = powershell.Invoke<T>();
+
+                if (result != null && result.Count > 0)
+                {
+                    output.AddRange(result);
+                }
+            }
+
+            return output;
+        }
+
+        /// <summary>
+        /// Generate a SHA256 Hash string from the originInput.
+        /// </summary>
+        /// <param name="originInput"></param>
+        /// <returns>The Sha256 hash, or empty if the input is only whtespace</returns>
+        private static string GenerateSha256HashString(string originInput)
+        {
+            if (string.IsNullOrWhiteSpace(originInput))
+            {
+                return string.Empty;
+            }
+
+            string result = string.Empty;
+            try
+            {
+                using (var sha256 = new SHA256CryptoServiceProvider())
+                {
+                    var bytes = sha256.ComputeHash(Encoding.UTF8.GetBytes(originInput));
+                    result = BitConverter.ToString(bytes);
+                }
+            }
+            catch
+            {
+                // do not throw if CryptoProvider is not provided
+            }
+
+            return result;
+        }
+
+        /// <summary>
+        /// Get the MAC address of the default NIC, or null if none can be found
+        /// </summary>
+        /// <returns>The MAC address of the defautl nic, or null if noen is found</returns>
+        private static string GetMACAddress()
+        {
+            return NetworkInterface.GetAllNetworkInterfaces()?
+                                    .FirstOrDefault(nic => nic != null &&
+                                                           nic.OperationalStatus == OperationalStatus.Up &&
+                                                           !string.IsNullOrWhiteSpace(nic.GetPhysicalAddress()?.ToString()))?
+                                    .GetPhysicalAddress()?.ToString();
+        }
+    }
+}

tools/Az.Tools.Predictor/Az.Tools.Predictor/AzPredictor.cs

@@ -57,6 +57,7 @@ internal sealed class AzPredictor : ICommandPredictor
         private readonly IAzPredictorService _service;
         private readonly ITelemetryClient _telemetryClient;
         private readonly Settings _settings;
+        private readonly IAzContext _azContext;
 
         private Queue<string> _lastTwoMaskedCommands = new Queue<string>(AzPredictorConstants.CommandHistoryCountToProcess);
 
@@ -66,16 +67,21 @@ internal sealed class AzPredictor : ICommandPredictor
         /// <param name="service">The service that provides the suggestion</param>
         /// <param name="telemetryClient">The client to collect telemetry</param>
         /// <param name="settings">The settings of the service</param>
-        public AzPredictor(IAzPredictorService service, ITelemetryClient telemetryClient, Settings settings)
+        /// <param name="azContext">The Az context which this module runs with</param>
+        public AzPredictor(IAzPredictorService service, ITelemetryClient telemetryClient, Settings settings, IAzContext azContext)
         {
             this._service = service;
             this._telemetryClient = telemetryClient;
             this._settings = settings;
+            this._azContext = azContext;
         }
 
         /// <inhericdoc />
         public void StartEarlyProcessing(IReadOnlyList<string> history)
         {
+            // The context only changes when the user executes the corresponding command.
+            this._azContext?.UpdateContext();
+
             if (history.Count > 0)
             {
                 if (_lastTwoMaskedCommands.Any())
@@ -262,9 +268,11 @@ public class PredictorInitializer : IModuleAssemblyInitializer
         public void OnImport()
         {
             var settings = Settings.GetSettings();
-            var telemetryClient = new AzPredictorTelemetryClient();
-            var azPredictorService = new AzPredictorService(settings.ServiceUri, telemetryClient);
-            var predictor = new AzPredictor(azPredictorService, telemetryClient, settings);
+            var azContext = new AzContext();
+            azContext.UpdateContext();
+            var telemetryClient = new AzPredictorTelemetryClient(azContext);
+            var azPredictorService = new AzPredictorService(settings.ServiceUri, telemetryClient, azContext);
+            var predictor = new AzPredictor(azPredictorService, telemetryClient, settings, azContext);
             SubsystemManager.RegisterSubsystem<ICommandPredictor, AzPredictor>(predictor);
         }
     }

tools/Az.Tools.Predictor/Az.Tools.Predictor/AzPredictorService.cs

@@ -19,6 +19,7 @@
 using System.Linq;
 using System.Management.Automation.Language;
 using System.Net.Http;
+using System.Net.Http.Headers;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
@@ -38,7 +39,7 @@ public sealed class RequestContext
                 public string CorrelationId { get; set; } = Guid.Empty.ToString();
                 public string SessionId { get; set; } = Guid.Empty.ToString();
                 public string SubscriptionId { get; set; } = Guid.Empty.ToString();
-                public Version VersionNumber{ get; set; } = new Version(1, 0);
+                public Version VersionNumber{ get; set; } = new Version(0, 0);
             }
 
             public string History { get; set; }
@@ -48,7 +49,8 @@ public sealed class RequestContext
             public PredictionRequestBody(string command) => this.History = command;
         };
 
-        private static readonly HttpClient _client = new HttpClient();
+        private const string ThrottleByIdHeader = "X-UserId";
+        private readonly HttpClient _client;
         private readonly string _commandsEndpoint;
         private readonly string _predictionsEndpoint;
         private volatile Tuple<string, Predictor> _commandSuggestions; // The command and the prediction for that.
@@ -59,18 +61,24 @@ public sealed class RequestContext
         private ParameterValuePredictor _parameterValuePredictor = new ParameterValuePredictor();
 
         private readonly ITelemetryClient _telemetryClient;
+        private readonly IAzContext _azContext;
 
         /// <summary>
         /// The AzPredictor service interacts with the Aladdin service specified in serviceUri.
         /// At initialization, it requests a list of the popular commands.
         /// </summary>
         /// <param name="serviceUri">The URI of the Aladdin service.</param>
         /// <param name="telemetryClient">The telemetry client.</param>
-        public AzPredictorService(string serviceUri, ITelemetryClient telemetryClient)
+        /// <param name="azContext">The Az context which this module runs with</param>
+        public AzPredictorService(string serviceUri, ITelemetryClient telemetryClient, IAzContext azContext)
         {
             this._commandsEndpoint = serviceUri + AzPredictorConstants.CommandsEndpoint;
             this._predictionsEndpoint = serviceUri + AzPredictorConstants.PredictionsEndpoint;
             this._telemetryClient = telemetryClient;
+            this._azContext = azContext;
+
+            this._client = new HttpClient();
+            this._client.DefaultRequestHeaders?.Add(AzPredictorService.ThrottleByIdHeader, this._azContext.HashedUserId);
 
             RequestCommands();
         }
@@ -163,6 +171,8 @@ public IEnumerable<ValueTuple<string, PredictionSource>> GetSuggestion(Ast input
         /// <inheritdoc/>
         public virtual void RequestPredictions(IEnumerable<string> commands)
         {
+            AzPredictorService.ReplaceThrottleUserIdToHeader(this._client?.DefaultRequestHeaders, this._azContext.HashedUserId);
+
             // Even if it's called multiple times, we only need to keep the one for the latest command.
 
             this._predictionRequestCancellationSource?.Cancel();
@@ -182,6 +192,7 @@ public virtual void RequestPredictions(IEnumerable<string> commands)
                     {
                         SessionId = this._telemetryClient.SessionId,
                         CorrelationId = this._telemetryClient.CorrelationId,
+                        VersionNumber = this._azContext.AzVersion,
                     };
                     var requestBody = new PredictionRequestBody(localCommands)
                     {
@@ -222,7 +233,7 @@ protected virtual void RequestCommands()
             // We don't need to block on the task. We send the HTTP request and update commands and predictions list at the background.
             Task.Run(async () =>
                     {
-                        var httpResponseMessage = await AzPredictorService._client.GetAsync(this._commandsEndpoint);
+                        var httpResponseMessage = await this._client.GetAsync(this._commandsEndpoint);
 
                         var reply = await httpResponseMessage.Content.ReadAsStringAsync();
                         var commands_reply = JsonConvert.DeserializeObject<List<string>>(reply);
@@ -271,5 +282,22 @@ private static string GetCommandName(string commandLine)
         {
             return commandLine.Split(AzPredictorConstants.CommandParameterSeperator).First();
         }
+
+        private static void ReplaceThrottleUserIdToHeader(HttpRequestHeaders header, string value)
+        {
+            if (header != null)
+            {
+                lock (header)
+                {
+                    header.Remove(AzPredictorService.ThrottleByIdHeader);
+
+                    if (!string.IsNullOrWhiteSpace(value))
+                    {
+                        header.Add(AzPredictorService.ThrottleByIdHeader, value);
+                    }
+                }
+            }
+
+        }
     }
 }