Enable parameter injection for Azure PowerShell response

build.psm1

@@ -91,6 +91,43 @@ function Start-Build
     $app_csproj = GetProjectFile $app_dir
     dotnet publish $app_csproj -c $Configuration -o $app_out_dir -r $RID --sc
 
+    ## Move the 'Modules' folder to the appbase folder.
+    if ($LASTEXITCODE -eq 0) {
+        ## Remove the artifacts that are not for the current platform, to reduce size.
+        $otherPlatDir = Join-Path $app_out_dir 'runtimes' ($IsWindows ? 'unix' : 'win')
+        if (Test-Path $otherPlatDir -PathType Container) {
+            Remove-Item $otherPlatDir -Recurse -Force -ErrorAction Stop
+        }
+
+        ## Move the 'Modules' folder and if possible, remove the 'runtimes' folder all together afterward.
+        $platDir = Join-Path $app_out_dir 'runtimes' ($IsWindows ? 'win' : 'unix') 'lib'
+        if (Test-Path $platDir -PathType Container) {
+            $moduleDir = Get-ChildItem $platDir -Directory -Include 'Modules' -Recurse
+            if ($moduleDir) {
+                ## Remove the existing 'Modules' folder if it already exists.
+                $target = Join-Path $app_out_dir 'Modules'
+                if (Test-Path $target -PathType Container) {
+                    Remove-Item $target -Recurse -Force -ErrorAction Stop
+                }
+
+                ## Move 'Modules' folder.
+                Move-Item $moduleDir.FullName $app_out_dir -Force -ErrorAction Stop
+
+                ## Remove the 'runtimes' folder if possible.
+                $parent = $moduleDir.Parent
+                while ($parent.FullName -ne $app_out_dir) {
+                    $files = Get-ChildItem $parent.FullName -File -Recurse
+                    if (-not $files) {
+                        Remove-Item $parent.FullName -Recurse -Force -ErrorAction Stop
+                        $parent = $parent.Parent
+                    } else {
+                        break
+                    }
+                }
+            }
+        }
+    }
+
     if ($LASTEXITCODE -eq 0 -and $AgentToInclude -contains 'openai-gpt') {
         Write-Host "`n[Build the OpenAI agent ...]`n" -ForegroundColor Green
         $openai_csproj = GetProjectFile $openai_agent_dir

shell/AIShell.App/AIShell.App.csproj

@@ -13,6 +13,7 @@
   <ItemGroup>
     <ProjectReference Include="..\AIShell.Kernel\AIShell.Kernel.csproj" />
     <PackageReference Include="System.CommandLine" Version="2.0.0-beta4.22272.1" />
+    <PackageReference Include="Microsoft.PowerShell.SDK" Version="7.4.7" />
   </ItemGroup>
 
 </Project>

shell/AIShell.Kernel/Utility/LoadContext.cs

@@ -1,11 +1,17 @@
-using System.Reflection;
+using System.Collections.Concurrent;
+using System.Reflection;
+using System.Runtime.InteropServices;
 using System.Runtime.Loader;
 
 namespace AIShell.Kernel;
 
 internal class AgentAssemblyLoadContext : AssemblyLoadContext
 {
     private readonly string _dependencyDir;
+    private readonly string _nativeLibExt;
+    private readonly List<string> _runtimeLibDir;
+    private readonly List<string> _runtimeNativeDir;
+    private readonly ConcurrentDictionary<string, Assembly> _cache;
 
     internal AgentAssemblyLoadContext(string name, string dependencyDir)
         : base($"{name.Replace(' ', '.')}-ALC", isCollectible: false)
@@ -17,20 +23,112 @@ internal AgentAssemblyLoadContext(string name, string dependencyDir)
 
         // Save the full path to the dependencies directory when creating the context.
         _dependencyDir = dependencyDir;
+        _runtimeLibDir = [];
+        _runtimeNativeDir = [];
+        _cache = [];
+
+        if (OperatingSystem.IsWindows())
+        {
+            _nativeLibExt = ".dll";
+            AddToList(_runtimeLibDir, Path.Combine(dependencyDir, "runtimes", "win", "lib"));
+            AddToList(_runtimeNativeDir, Path.Combine(dependencyDir, "runtimes", "win", "native"));
+        }
+        else if (OperatingSystem.IsLinux())
+        {
+            _nativeLibExt = ".so";
+            AddToList(_runtimeLibDir, Path.Combine(dependencyDir, "runtimes", "unix", "lib"));
+            AddToList(_runtimeLibDir, Path.Combine(dependencyDir, "runtimes", "linux", "lib"));
+
+            AddToList(_runtimeNativeDir, Path.Combine(dependencyDir, "runtimes", "unix", "native"));
+            AddToList(_runtimeNativeDir, Path.Combine(dependencyDir, "runtimes", "linux", "native"));
+        }
+        else if (OperatingSystem.IsMacOS())
+        {
+            _nativeLibExt = ".dylib";
+            AddToList(_runtimeLibDir, Path.Combine(dependencyDir, "runtimes", "unix", "lib"));
+            AddToList(_runtimeLibDir, Path.Combine(dependencyDir, "runtimes", "osx", "lib"));
+
+            AddToList(_runtimeNativeDir, Path.Combine(dependencyDir, "runtimes", "unix", "native"));
+            AddToList(_runtimeNativeDir, Path.Combine(dependencyDir, "runtimes", "osx", "native"));
+        }
+
+        AddToList(_runtimeLibDir, Path.Combine(dependencyDir, "runtimes", RuntimeInformation.RuntimeIdentifier, "lib"));
+        AddToList(_runtimeNativeDir, Path.Combine(dependencyDir, "runtimes", RuntimeInformation.RuntimeIdentifier, "native"));
+
+        ResolvingUnmanagedDll += ResolveUnmanagedDll;
     }
 
     protected override Assembly Load(AssemblyName assemblyName)
     {
-        // Create a path to the assembly in the dependencies directory.
-        string path = Path.Combine(_dependencyDir, $"{assemblyName.Name}.dll");
+        if (_cache.TryGetValue(assemblyName.Name, out Assembly assembly))
+        {
+            return assembly;
+        }
+
+        lock (this)
+        {
+            if (_cache.TryGetValue(assemblyName.Name, out assembly))
+            {
+                return assembly;
+            }
+
+            // Create a path to the assembly in the dependencies directory.
+            string assemblyFile = $"{assemblyName.Name}.dll";
+            string path = Path.Combine(_dependencyDir, assemblyFile);
+
+            if (File.Exists(path))
+            {
+                // If the assembly exists in our dependency directory, then load it into this load context.
+                assembly = LoadFromAssemblyPath(path);
+            }
+            else
+            {
+                foreach (string dir in _runtimeLibDir)
+                {
+                    IEnumerable<string> result = Directory.EnumerateFiles(dir, assemblyFile, SearchOption.AllDirectories);
+                    path = result.FirstOrDefault();
+
+                    if (path is not null)
+                    {
+                        assembly = LoadFromAssemblyPath(path);
+                        break;
+                    }
+                }
+            }
+
+            // Add the probing result to cache, regardless of whether we found it.
+            // If we didn't find it, we will add 'null' to the cache so that we don't probe
+            // again in case another loading request comes for the same assembly.
+            _cache.TryAdd(assemblyName.Name, assembly);
+
+            // Return the assembly if we found it, or return 'null' otherwise to depend on the default load context to resolve the request.
+            return assembly;
+        }
+    }
+
+    private nint ResolveUnmanagedDll(Assembly assembly, string libraryName)
+    {
+        string libraryFile = $"{libraryName}{_nativeLibExt}";
 
-        if (File.Exists(path))
+        foreach (string dir in _runtimeNativeDir)
         {
-            // If the assembly exists in our dependency directory, then load it into this load context.
-            return LoadFromAssemblyPath(path);
+            IEnumerable<string> result = Directory.EnumerateFiles(dir, libraryFile, SearchOption.AllDirectories);
+            string path = result.FirstOrDefault();
+
+            if (path is not null)
+            {
+                return NativeLibrary.Load(path);
+            }
         }
 
-        // Otherwise we will depend on the default load context to resolve the request.
-        return null;
+        return nint.Zero;
+    }
+
+    private static void AddToList(List<string> depList, string dirPath)
+    {
+        if (Directory.Exists(dirPath))
+        {
+            depList.Add(dirPath);
+        }
     }
 }

shell/agents/Microsoft.Azure.Agent/AzureAgent.cs

@@ -27,7 +27,7 @@ public sealed class AzureAgent : ILLMAgent
         2. DO NOT include the command for creating a new resource group unless the query explicitly asks for it. Otherwise, assume a resource group already exists.
         3. DO NOT include an additional example with made-up values unless it provides additional context or value beyond the initial command.
         4. DO NOT use the line continuation operator (backslash `\` in Bash) in the generated commands.
-        5. Always represent a placeholder in the form of `<placeholder-name>`.
+        5. Always represent a placeholder in the form of `<placeholder-name>` and enclose it within double quotes.
         6. Always use the consistent placeholder names across all your responses. For example, `<resourceGroupName>` should be used for all the places where a resource group name value is needed.
         7. When the commands contain placeholders, the placeholders should be summarized in markdown bullet points at the end of the response in the same order as they appear in the commands, following this format:
            ```
@@ -260,9 +260,9 @@ public async Task<bool> ChatAsync(string input, IShell shell)
                 {
                     // Process CLI handler response specially to support parameter injection.
                     ResponseData data = null;
-                    if (_copilotResponse.TopicName == CopilotActivity.CLIHandlerTopic)
+                    if (_copilotResponse.TopicName is CopilotActivity.CLIHandlerTopic or CopilotActivity.PSHandlerTopic)
                     {
-                        data = ParseCLIHandlerResponse(shell);
+                        data = ParseCodeResponse(shell);
                     }
 
                     if (data?.PlaceholderSet is not null)
@@ -349,7 +349,7 @@ public async Task<bool> ChatAsync(string input, IShell shell)
         return true;
     }
 
-    private ResponseData ParseCLIHandlerResponse(IShell shell)
+    private ResponseData ParseCodeResponse(IShell shell)
     {
         string text = _copilotResponse.Text;
         List<CodeBlock> codeBlocks = shell.ExtractCodeBlocks(text, out List<SourceInfo> sourceInfos);
@@ -402,6 +402,7 @@ private ResponseData ParseCLIHandlerResponse(IShell shell)
         ResponseData data = new() {
             Text = text,
             CommandSet = commands,
+            TopicName = _copilotResponse.TopicName,
             PlaceholderSet = placeholders,
             Locale = _copilotResponse.Locale,
         };

shell/agents/Microsoft.Azure.Agent/Command.cs

@@ -33,9 +33,19 @@ private static string SyntaxHighlightAzCommand(string command, string parameter,
         const string vtReset = "\x1b[0m";
 
         StringBuilder cStr = new(capacity: command.Length + parameter.Length + placeholder.Length + 50);
-        cStr.Append(vtItalic)
-            .Append(vtCommand).Append("az").Append(vtFgDefault).Append(command.AsSpan(2)).Append(' ')
-            .Append(vtParameter).Append(parameter).Append(vtFgDefault).Append(' ')
+        cStr.Append(vtItalic);
+
+        int index = command.IndexOf(' ');
+        if (index is -1)
+        {
+            cStr.Append(vtCommand).Append(command).Append(vtFgDefault).Append(' ');
+        }
+        else
+        {
+            cStr.Append(vtCommand).Append(command.AsSpan(0, index)).Append(vtFgDefault).Append(command.AsSpan(index)).Append(' ');
+        }
+
+        cStr.Append(vtParameter).Append(parameter).Append(vtFgDefault).Append(' ')
             .Append(vtVariable).Append(placeholder).Append(vtFgDefault)
             .Append(vtReset);
 
@@ -125,15 +135,9 @@ private void ReplaceAction()
 
                 // Prompt for argument without printing captions again.
                 string value = host.PromptForArgument(argInfo, printCaption: false);
+                value = value?.Trim();
                 if (!string.IsNullOrEmpty(value))
                 {
-                    // Add quotes for the value if needed.
-                    value = value.Trim();
-                    if (value.StartsWith('-') || value.Contains(' ') || value.Contains('|'))
-                    {
-                        value = $"\"{value}\"";
-                    }
-
                     _values.Add(item.Name, value);
                     _agent.SaveUserValue(item.Name, value);