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>
@@ -157,7 +158,8 @@ public void VerifySuggestionOnIncompleteCommand()
             var localAzPredictor = new AzPredictor(this._service, this._telemetryClient, new Settings()
             {
                 SuggestionCount = 7,
-            });
+            },
+            null);
 
             var userInput = "New-AzResourceGroup -Name 'ResourceGroup01' -Location 'Central US' -WhatIf -";
             var expected = "New-AzResourceGroup -Name 'ResourceGroup01' -Location 'Central US' -WhatIf -Verbose ***";

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,186 @@
+// ----------------------------------------------------------------------------------
+//
+// 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
+    {
+        private static readonly Version DefaultVersion = new Version("0.0.0.0");
+
+        /// <inheritdoc/>
+        public string UserId { get; private set; } = string.Empty;
+
+        private string _macAddress;
+        /// <inheritdoc/>
+        public string MacAddress
+        {
+            get
+            {
+                if (_macAddress == null)
+                {
+                    _macAddress = string.Empty;
+
+                    var macAddress = GetMACAddress();
+                    if (!string.IsNullOrWhiteSpace(macAddress))
+                    {
+                        _macAddress = GenerateSha256HashString(macAddress)?.Replace("-", string.Empty).ToLowerInvariant();
+                    }
+                }
+
+                return _macAddress;
+            }
+        }
+
+        /// <inheritdoc/>
+        public string OSVersion
+        {
+            get
+            {
+                return Environment.OSVersion.ToString();
+            }
+        }
+
+        private Version _powerShellVersion;
+        /// <inheritdoc/>
+        public Version PowerShellVersion
+        {
+            get
+            {
+                if (_powerShellVersion == null)
+                {
+                    var outputs = AzContext.ExecuteScript<Version>("(Get-Host).Version");
+
+                    _powerShellVersion = outputs.FirstOrDefault();
+                }
+
+                return _powerShellVersion ?? AzContext.DefaultVersion;
+            }
+        }
+
+        private Version _moduleVersion;
+        /// <inheritdoc/>
+        public Version ModuleVersion
+        {
+            get
+            {
+                if (_moduleVersion == null)
+                {
+                    _moduleVersion = this.GetType().Assembly.GetName().Version;
+                }
+
+                return _moduleVersion ?? AzContext.DefaultVersion;
+            }
+        }
+
+        /// <inheritdoc/>
+        public void UpdateContext()
+        {
+            UserId = GenerateSha256HashString(GetUserAccountId());
+        }
+
+        /// <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 whitespace</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 none 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);
 
@@ -69,16 +70,20 @@ 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();
             lock (_userAcceptedAndSuggestion)
             {
                 _userAcceptedAndSuggestion.Clear();
@@ -330,9 +335,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;
@@ -56,7 +57,8 @@ private sealed class CommandRequestContext
             public Version VersionNumber{ get; set; } = new Version(0, 0);
         }
 
-        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.
@@ -67,18 +69,24 @@ private sealed class CommandRequestContext
         private readonly 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}?clientType={AzPredictorService.ClientType}&context={JsonConvert.SerializeObject(new CommandRequestContext())}";
             this._predictionsEndpoint = serviceUri + AzPredictorConstants.PredictionsEndpoint;
             this._telemetryClient = telemetryClient;
+            this._azContext = azContext;
+
+            this._client = new HttpClient();
+            this._client.DefaultRequestHeaders?.Add(AzPredictorService.ThrottleByIdHeader, this._azContext.UserId);
 
             RequestCommands();
         }
@@ -173,6 +181,7 @@ public IEnumerable<ValueTuple<string, IList<Tuple<string, string>>, PredictionSo
         /// <inheritdoc/>
         public virtual void RequestPredictions(IEnumerable<string> commands)
         {
+            AzPredictorService.ReplaceThrottleUserIdToHeader(this._client?.DefaultRequestHeaders, this._azContext.UserId);
             var localCommands= string.Join(AzPredictorConstants.CommandConcatenator, commands);
             this._telemetryClient.OnRequestPrediction(localCommands);
 
@@ -241,7 +250,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);
@@ -290,5 +299,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);
+                    }
+                }
+            }
+
+        }
     }
 }