Add support to Azure PowerShell login credential (#329)

PowerShell · Jan 21, 2025 · 1c4a2e8 · 1c4a2e8
1 parent d16ac5a
commit 1c4a2e8
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 5 deletions.
diff --git a/shell/agents/Microsoft.Azure.Agent/AzureAgent.cs b/shell/agents/Microsoft.Azure.Agent/AzureAgent.cs
@@ -194,7 +194,7 @@ public async Task RefreshChatAsync(IShell shell, bool force)
             if (inner is CredentialUnavailableException)
             {
                 host.WriteErrorLine($"Failed to start a chat session: Access token not available.");
-                host.WriteErrorLine($"The '{Name}' agent depends on the Azure CLI credential to acquire access token. Please run 'az login' from a command-line shell to setup account. Once you've successfully logged in, run '/refresh' to start a new conversation.");
+                host.WriteErrorLine($"The '{Name}' agent depends on the Azure CLI credential or Azure PowerShell credential to acquire access token. Please run 'az login' or 'Connect-AzAccount' from a command-line shell to setup account. Once you've successfully logged in, run '/refresh' to start a new conversation.");
                 return;
             }
 

diff --git a/shell/agents/Microsoft.Azure.Agent/ChatSession.cs b/shell/agents/Microsoft.Azure.Agent/ChatSession.cs
@@ -115,8 +115,8 @@ private async Task<string> SetupNewChat(IStatusContext context, CancellationToke
     {
         try
         {
-            context.Status("Get Azure CLI login token ...");
-            // Get an access token from the AzCLI login, using the specific audience guid.
+            context.Status("Get Azure login token ...");
+            // Get an access token from the Azure CLI/PowerShell login, using the specific audience guid.
             await _accessToken.CreateOrRenewTokenAsync(cancellationToken);
 
             context.Status("Check Copilot authorization ...");

diff --git a/shell/agents/Microsoft.Azure.Agent/Schema.cs b/shell/agents/Microsoft.Azure.Agent/Schema.cs
@@ -386,8 +386,10 @@ internal async Task CreateOrRenewTokenAsync(CancellationToken cancellationToken)
 
             if (needRefresh)
             {
-                _accessToken = await new AzureCliCredential()
-                    .GetTokenAsync(_tokenContext, cancellationToken);
+                _accessToken = await new ChainedTokenCredential(
+                    new AzureCliCredential(),
+                    new AzurePowerShellCredential()
+                ).GetTokenAsync(_tokenContext, cancellationToken);
             }
         }
         catch (Exception e) when (e is not OperationCanceledException)