Add PwshHandler for native PowerShell 7.0+ Tasks

src/Agent.Worker/Handlers/HandlerFactory.cs

@@ -76,12 +76,16 @@ public IHandler Create(
             }
             else if (data is PowerShell3HandlerData)
             {
-                #pragma warning disable CA1416 // PowerShell handlers are Windows only
-                // PowerShell3.
+                #pragma warning disable CA1416 // PowerShell3 handler is Windows only
                 handler = HostContext.CreateService<IPowerShell3Handler>();
                 (handler as IPowerShell3Handler).Data = data as PowerShell3HandlerData;
                 #pragma warning restore CA1416
             }
+            else if (data is PwshHandlerData)
+            {
+                handler = HostContext.CreateService<IPwshHandler>();
+                (handler as IPwshHandler).Data = data as PwshHandlerData;
+            }
             else if (data is PowerShellExeHandlerData)
             {
                 #pragma warning disable CA1416 // PowerShell handlers are Windows only
@@ -169,4 +173,4 @@ private List<Guid> getTaskExceptionList()
             return exceptionList;
         }
     }
-}
+}

src/Agent.Worker/Handlers/PwshHandler.cs

@@ -0,0 +1,143 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using Agent.Sdk;
+using Agent.Sdk.Knob;
+using Microsoft.VisualStudio.Services.Agent.Util;
+using System;
+using System.IO;
+using System.Threading.Tasks;
+
+namespace Microsoft.VisualStudio.Services.Agent.Worker.Handlers
+{
+    [ServiceLocator(Default = typeof(PwshHandler))]
+    public interface IPwshHandler : IHandler
+    {
+        PwshHandlerData Data { get; set; }
+    }
+
+    public sealed class PwshHandler : Handler, IPwshHandler
+    {
+        public PwshHandlerData Data { get; set; }
+
+        public async Task RunAsync()
+        {
+            Trace.Entering();
+            ArgUtil.NotNull(Data, nameof(Data));
+            ArgUtil.NotNull(ExecutionContext, nameof(ExecutionContext));
+            ArgUtil.NotNull(Inputs, nameof(Inputs));
+            ArgUtil.Directory(TaskDirectory, nameof(TaskDirectory));
+
+            AddInputsToEnvironment();
+            AddEndpointsToEnvironment();
+            AddSecureFilesToEnvironment();
+            AddVariablesToEnvironment();
+            AddTaskVariablesToEnvironment();
+            AddPrependPathToEnvironment();
+            if (PlatformUtil.RunningOnWindows)
+            {
+                RemovePSModulePathFromEnvironment();
+            }
+
+            string scriptFile = ResolveScriptFile();
+            string scriptDirectory = Path.GetDirectoryName(scriptFile);
+            string moduleFile = ResolveModuleFile(scriptDirectory);
+            string pwshArgs = BuildPwshArguments(moduleFile, scriptFile);
+            string pwsh = ResolvePwshExecutable();
+
+            StepHost.OutputDataReceived += OnDataReceived;
+            StepHost.ErrorDataReceived += OnDataReceived;
+
+            var sigintTimeout = TimeSpan.FromMilliseconds(AgentKnobs.ProccessSigintTimeout.GetValue(ExecutionContext).AsInt());
+            var sigtermTimeout = TimeSpan.FromMilliseconds(AgentKnobs.ProccessSigtermTimeout.GetValue(ExecutionContext).AsInt());
+            var useGracefulShutdown = AgentKnobs.UseGracefulProcessShutdown.GetValue(ExecutionContext).AsBoolean();
+
+            try
+            {
+                await StepHost.ExecuteAsync(
+                    workingDirectory: StepHost.ResolvePathForStepHost(scriptDirectory),
+                    fileName: pwsh,
+                    arguments: pwshArgs,
+                    environment: Environment,
+                    requireExitCodeZero: true,
+                    outputEncoding: null,
+                    killProcessOnCancel: false,
+                    inheritConsoleHandler: !ExecutionContext.Variables.Retain_Default_Encoding,
+                    continueAfterCancelProcessTreeKillAttempt: _continueAfterCancelProcessTreeKillAttempt,
+                    sigintTimeout: sigintTimeout,
+                    sigtermTimeout: sigtermTimeout,
+                    useGracefulShutdown: useGracefulShutdown,
+                    cancellationToken: ExecutionContext.CancellationToken);
+            }
+            finally
+            {
+                StepHost.OutputDataReceived -= OnDataReceived;
+                StepHost.ErrorDataReceived -= OnDataReceived;
+            }
+        }
+
+        private void OnDataReceived(object sender, ProcessDataReceivedEventArgs e)
+        {
+            if (!CommandManager.TryProcessCommand(ExecutionContext, e.Data))
+            {
+                ExecutionContext.Output(e.Data);
+            }
+        }
+
+        private string ResolveScriptFile()
+        {
+            ArgUtil.NotNullOrEmpty(Data.Target, nameof(Data.Target));
+            string scriptFile = Path.Combine(TaskDirectory, Data.Target);
+            ArgUtil.File(scriptFile, nameof(scriptFile));
+            return scriptFile;
+        }
+
+        private string ResolveModuleFile(string scriptDirectory)
+        {
+            string moduleFile = Path.Combine(scriptDirectory, "ps_modules", "VstsTaskSdk", "VstsTaskSdk.psd1");
+            ArgUtil.File(moduleFile, nameof(moduleFile));
+            return moduleFile;
+        }
+
+        private string BuildPwshArguments(string moduleFile, string scriptFile)
+        {
+            if (AgentKnobs.UsePSScriptWrapper.GetValue(ExecutionContext).AsBoolean())
+            {
+                return BuildWrapperArguments(moduleFile, scriptFile);
+            }
+
+            return BuildDirectInvocationArguments(moduleFile, scriptFile);
+        }
+
+        private string BuildWrapperArguments(string moduleFile, string scriptFile)
+        {
+            return StringUtil.Format(
+                @"-NoLogo -NoProfile -ExecutionPolicy Unrestricted -Command ""{3}"" -VstsSdkPath {0} -DebugOption {1} -ScriptBlockString ""{2}""",
+                StepHost.ResolvePathForStepHost(moduleFile).Replace("'", "''"),
+                ExecutionContext.Variables.System_Debug == true ? "Continue" : "SilentlyContinue",
+                StepHost.ResolvePathForStepHost(scriptFile).Replace("'", "''''"),
+                Path.Combine(HostContext.GetDirectory(WellKnownDirectory.Bin), "powershell", "Start-AzpTask.ps1"));
+        }
+
+        private string BuildDirectInvocationArguments(string moduleFile, string scriptFile)
+        {
+            return StringUtil.Format(
+                @"-NoLogo -NoProfile -NonInteractive -ExecutionPolicy Unrestricted -Command "". ([scriptblock]::Create('if ([Console]::InputEncoding -is [Text.UTF8Encoding] -and [Console]::InputEncoding.GetPreamble().Length -ne 0) {{ [Console]::InputEncoding = New-Object Text.UTF8Encoding $false }}')) 2>&1 | ForEach-Object {{ Write-Verbose $_.Exception.Message -Verbose }} ; Import-Module -Name '{0}' -ArgumentList @{{ NonInteractive = $true }} -ErrorAction Stop ; $VerbosePreference = '{1}' ; $DebugPreference = '{1}' ; Invoke-VstsTaskScript -ScriptBlock ([scriptblock]::Create('. ''{2}'''))""",
+                StepHost.ResolvePathForStepHost(moduleFile).Replace("'", "''"),
+                ExecutionContext.Variables.System_Debug == true ? "Continue" : "SilentlyContinue",
+                StepHost.ResolvePathForStepHost(scriptFile).Replace("'", "''''"));
+        }
+
+        private string ResolvePwshExecutable()
+        {
+            string pwsh = "pwsh";
+            if (StepHost is DefaultStepHost)
+            {
+                pwsh = HostContext.GetService<IPwshExeUtil>().GetPath();
+            }
+
+            ArgUtil.NotNullOrEmpty(pwsh, nameof(pwsh));
+            return pwsh;
+        }
+    }
+}

src/Agent.Worker/TaskManager.cs

@@ -616,6 +616,7 @@ public sealed class ExecutionData
         private Node16HandlerData _node16;
         private Node20_1HandlerData _node20_1;
         private Node24HandlerData _node24;
+        private PwshHandlerData _pwsh;
         private PowerShellHandlerData _powerShell;
         private PowerShell3HandlerData _powerShell3;
         private PowerShellExeHandlerData _powerShellExe;
@@ -713,6 +714,21 @@ public Node24HandlerData Node24
             }
         }
 
+        [JsonProperty(NullValueHandling = NullValueHandling.Ignore)]
+        public PwshHandlerData Pwsh
+        {
+            get
+            {
+                return _pwsh;
+            }
+
+            set
+            {
+                _pwsh = value;
+                Add(value);
+            }
+        }
+
         [JsonProperty(NullValueHandling = NullValueHandling.Ignore)]
         public PowerShellHandlerData PowerShell
         {
@@ -916,6 +932,11 @@ public sealed class PowerShell3HandlerData : HandlerData
         public override int Priority => 106;
     }
 
+    public sealed class PwshHandlerData : HandlerData
+    {
+        public override int Priority => 106;
+    }
+
     public sealed class PowerShellHandlerData : HandlerData
     {
         public string ArgumentFormat

src/Agent.Worker/TaskRunner.cs

@@ -202,7 +202,7 @@ private async Task RunAsyncInternal()
                         runtimeVariables = new Variables(HostContext, variableCopy, out expansionWarnings);
                         expansionWarnings?.ForEach(x => ExecutionContext.Warning(x));
                     }
-                    else if (handlerData is BaseNodeHandlerData || handlerData is PowerShell3HandlerData)
+                    else if (handlerData is BaseNodeHandlerData || handlerData is PowerShell3HandlerData || handlerData is PwshHandlerData)
                     {
                         // Only the node, node10, and powershell3 handlers support running inside container.
                         // Make sure required container is already created.
@@ -580,7 +580,7 @@ public HandlerData GetHandlerData(IExecutionContext ExecutionContext, ExecutionD
                 targetOS = stepTarget.ExecutionOS;
                 if (stepTarget is ContainerInfo)
                 {
-                    if ((currentExecution.All.Any(x => x is PowerShell3HandlerData)) &&
+                    if ((currentExecution.All.Any(x => x is PowerShell3HandlerData || x is PwshHandlerData)) &&
                         (currentExecution.All.Any(x => x is BaseNodeHandlerData)))
                     {
                         Trace.Info($"Since we are targeting a container, we will prefer a node handler if one is available");
@@ -590,7 +590,7 @@ public HandlerData GetHandlerData(IExecutionContext ExecutionContext, ExecutionD
             }
             Trace.Info($"Get handler data for target platform {targetOS.ToString()}");
             return currentExecution.All
-                    .OrderBy(x => !(x.PreferredOnPlatform(targetOS) && (preferPowershellHandler || !(x is PowerShell3HandlerData)))) // Sort true to false.
+                    .OrderBy(x => !(x.PreferredOnPlatform(targetOS) && (preferPowershellHandler || !(x is PowerShell3HandlerData || x is PwshHandlerData)))) // Sort true to false.
                     .ThenBy(x => x.Priority)
                     .FirstOrDefault();
         }
@@ -742,9 +742,9 @@ private bool IsCorrelationIdRequired(IHandler handler, Definition task)
                     Trace.Info($"Node SDK version: {nodeSdkVer}. Correlation ID is required: {isIdRequired}.");
                 }
             }
-            else if (handler is IPowerShell3Handler)
+            else if (handler is IPowerShell3Handler || handler is IPwshHandler)
             {
-                Trace.Info("Current handler is PowerShell3. Trying to determing the SDK version.");
+                Trace.Info("Current handler is PowerShell. Trying to determing the SDK version.");
                 var psSdkVer = task.GetPowerShellSDKVersion();
                 if (psSdkVer == null)
                 {
@@ -766,4 +766,4 @@ private bool IsCorrelationIdRequired(IHandler handler, Definition task)
             return isIdRequired;
         }
     }
-}
+}

src/Microsoft.VisualStudio.Services.Agent/Capabilities/PwshCapabilitiesProvider.cs

@@ -0,0 +1,36 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using Microsoft.VisualStudio.Services.Agent.Util;
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.VisualStudio.Services.Agent.Capabilities
+{
+    public sealed class PwshCapabilitiesProvider : AgentService, ICapabilitiesProvider
+    {
+        public Type ExtensionType => typeof(ICapabilitiesProvider);
+
+        public int Order => 2;
+
+        public Task<List<Capability>> GetCapabilitiesAsync(AgentSettings settings, CancellationToken cancellationToken)
+        {
+            Trace.Entering();
+            var capabilities = new List<Capability>();
+
+            try
+            {
+                string pwsh = HostContext.GetService<IPwshExeUtil>().GetPath();
+                capabilities.Add(new Capability("Pwsh", pwsh));
+            }
+            catch (Exception ex)
+            {
+                Trace.Info($"Pwsh capability not detected: {ex.Message}");
+            }
+
+            return Task.FromResult(capabilities);
+        }
+    }
+}

src/Microsoft.VisualStudio.Services.Agent/ExtensionManager.cs

@@ -46,6 +46,7 @@ private List<IExtension> LoadExtensions<T>() where T : class, IExtension
                 case "Microsoft.VisualStudio.Services.Agent.Capabilities.ICapabilitiesProvider":
                     Add<T>(extensions, "Microsoft.VisualStudio.Services.Agent.Capabilities.AgentCapabilitiesProvider, Microsoft.VisualStudio.Services.Agent");
                     Add<T>(extensions, "Microsoft.VisualStudio.Services.Agent.Capabilities.EnvironmentCapabilitiesProvider, Microsoft.VisualStudio.Services.Agent");
+                    Add<T>(extensions, "Microsoft.VisualStudio.Services.Agent.Capabilities.PwshCapabilitiesProvider, Microsoft.VisualStudio.Services.Agent");
                     if (PlatformUtil.RunningOnLinux || PlatformUtil.RunningOnMacOS)
                     {
                         Add<T>(extensions, "Microsoft.VisualStudio.Services.Agent.Capabilities.NixCapabilitiesProvider, Microsoft.VisualStudio.Services.Agent");
@@ -145,4 +146,4 @@ private void Add<T>(List<IExtension> extensions, string assemblyQualifiedName) w
             extensions.Add(extension);
         }
     }
-}
+}

src/Microsoft.VisualStudio.Services.Agent/Util/PwshExeUtil.cs

@@ -0,0 +1,89 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using Agent.Sdk;
+using Agent.Sdk.Util;
+using System;
+using System.Diagnostics;
+using System.Text.RegularExpressions;
+
+namespace Microsoft.VisualStudio.Services.Agent.Util
+{
+    [ServiceLocator(Default = typeof(PwshExeUtil))]
+    public interface IPwshExeUtil : IAgentService
+    {
+        string GetPath();
+    }
+
+    public sealed class PwshExeUtil : AgentService, IPwshExeUtil
+    {
+        private static readonly Version MinimumVersion = new Version(7, 0);
+
+        public string GetPath()
+        {
+            Trace.Entering();
+
+            string commandName = PlatformUtil.RunningOnWindows ? "pwsh.exe" : "pwsh";
+            string pwshPath = WhichUtil.Which(commandName, trace: Trace);
+            if (string.IsNullOrEmpty(pwshPath))
+            {
+                throw new InvalidOperationException(StringUtil.Loc("FileNotFound", commandName));
+            }
+
+            Version version = GetVersion(pwshPath);
+            if (version < MinimumVersion)
+            {
+                throw new InvalidOperationException($"A compatible version of pwsh was not found. Minimum required version is {MinimumVersion}.");
+            }
+
+            return pwshPath;
+        }
+
+        private Version GetVersion(string pwshPath)
+        {
+            ArgUtil.NotNullOrEmpty(pwshPath, nameof(pwshPath));
+
+            string output = string.Empty;
+            using (var process = new Process())
+            {
+                process.StartInfo = new ProcessStartInfo
+                {
+                    FileName = pwshPath,
+                    Arguments = "-NoLogo -NoProfile -NonInteractive -Command \"$PSVersionTable.PSVersion.ToString()\"",
+                    RedirectStandardOutput = true,
+                    RedirectStandardError = true,
+                    UseShellExecute = false,
+                    CreateNoWindow = true,
+                };
+
+                if (!process.Start())
+                {
+                    throw new InvalidOperationException($"Unable to start '{pwshPath}'.");
+                }
+
+                output = process.StandardOutput.ReadToEnd();
+                string error = process.StandardError.ReadToEnd();
+                process.WaitForExit();
+
+                if (process.ExitCode != 0)
+                {
+                    throw new InvalidOperationException($"Unable to determine pwsh version from '{pwshPath}'. {error}".Trim());
+                }
+            }
+
+            string versionString = output?.Trim();
+            if (Version.TryParse(versionString, out Version version))
+            {
+                return version;
+            }
+
+            Match match = Regex.Match(versionString ?? string.Empty, @"\d+\.\d+(\.\d+)?(\.\d+)?", RegexOptions.CultureInvariant);
+            if (match.Success && Version.TryParse(match.Value, out version))
+            {
+                return version;
+            }
+
+            throw new InvalidOperationException($"Unable to parse pwsh version '{versionString}'.");
+        }
+    }
+}