Refactor MCP token provider selection and fallback logic

src/Tests/Microsoft.Agents.A365.Tooling.Core.Tests/DevMcpTokenProviderTests.cs → src/Tests/Microsoft.Agents.A365.Tooling.Core.Tests/EnvMcpTokenProviderTests.cs

@@ -13,14 +13,14 @@
 namespace Microsoft.Agents.A365.Tooling.Core.Tests;
 
 /// <summary>
-/// Tests for <see cref="DevMcpTokenProvider"/> and <see cref="Utility.IsDevScenario"/>.
+/// Tests for <see cref="EnvMcpTokenProvider"/> and <see cref="Utility.IsDevScenario"/>.
 /// </summary>
-public class DevMcpTokenProviderTests
+public class EnvMcpTokenProviderTests
 {
     private static MCPServerConfig Server(string name) =>
         new() { mcpServerName = name, id = $"id-{name}", url = "http://test" };
 
-    private static DevMcpTokenProvider Provider(IConfiguration config) =>
+    private static EnvMcpTokenProvider Provider(IConfiguration config) =>
         new(config, Mock.Of<ILogger>());
 
     private static IConfiguration Config(params (string key, string value)[] entries)
@@ -91,16 +91,6 @@ public async Task GetTokenAsync_MultipleServers_SharedFallbackUsedForAll()
         t2.Should().Be("shared-token");
     }
 
-    // ─── Missing token → exception ────────────────────────────────────────────
-
-    [Fact]
-    public async Task GetTokenAsync_NeitherEnvVarSet_ThrowsInvalidOperationException()
-    {
-        var config = Config(); // empty — no BEARER_TOKEN_* or BEARER_TOKEN
-        Func<Task> act = () => Provider(config).GetTokenAsync(Server("mcp_MailTools"));
-        await act.Should().ThrowAsync<InvalidOperationException>()
-            .WithMessage("*mcp_MailTools*");
-    }
 
     [Fact]
     public async Task GetTokenAsync_PerServerVarWhitespaceOnly_FallsBackToSharedToken()
@@ -114,29 +104,6 @@ public async Task GetTokenAsync_PerServerVarWhitespaceOnly_FallsBackToSharedToke
         token.Should().Be("shared-token");
     }
 
-    [Fact]
-    public async Task GetTokenAsync_BothVarsWhitespaceOnly_ThrowsInvalidOperationException()
-    {
-        var config = Config(
-            ("BEARER_TOKEN_MCP_MAILTOOLS", "   "),
-            ("BEARER_TOKEN", "   "));
-
-        Func<Task> act = () => Provider(config).GetTokenAsync(Server("mcp_MailTools"));
-        await act.Should().ThrowAsync<InvalidOperationException>();
-    }
-
-    // ─── Error message quality ────────────────────────────────────────────────
-
-    [Fact]
-    public async Task GetTokenAsync_MissingToken_ErrorMessageContainsNormalizedKey()
-    {
-        var config = Config();
-        Func<Task> act = () => Provider(config).GetTokenAsync(Server("mcp_MailTools"));
-
-        var ex = await act.Should().ThrowAsync<InvalidOperationException>();
-        ex.Which.Message.Should().Contain("BEARER_TOKEN_MCP_MAILTOOLS");
-        ex.Which.Message.Should().Contain("BEARER_TOKEN");
-    }
 
     // ─── Cancellation ─────────────────────────────────────────────────────────
 
@@ -214,7 +181,7 @@ public void IsDevScenario_NoEnvironmentSet_ReturnsFalse()
     {
         // Unset environment must NOT default to Development so that hosts without an
         // explicit ASPNETCORE_ENVIRONMENT / DOTNET_ENVIRONMENT are not silently treated
-        // as dev (which would enable manifest discovery, DevMcpTokenProvider, and relaxed TLS).
+        // as dev (which would enable manifest discovery, EnvMcpTokenProvider, and relaxed TLS).
         var config = Config(); // nothing set
         Utility.IsDevScenario(config).Should().BeFalse();
     }

src/Tests/Microsoft.Agents.A365.Tooling.Extensions.AzureAIFoundry.Tests/McpToolRegistrationServiceTests.cs

@@ -93,12 +93,13 @@ private void SetupMocksForEmptyToolEnumeration(Action<ToolOptions>? captureToolO
             .Setup(x => x.EnumerateToolsFromServersAsync(
                 It.IsAny<string>(),
                 It.IsAny<string>(),
+                It.IsAny<TokenProviderCollection>(),
                 It.IsAny<ITurnContext>(),
                 It.IsAny<ToolOptions>()));
 
         if (captureToolOptions != null)
         {
-            setup.Callback<string, string, ITurnContext, ToolOptions>((_, _, _, options) => captureToolOptions(options));
+            setup.Callback<string, string, IMcpTokenProvider, ITurnContext, ToolOptions>((_, _, _, _, options) => captureToolOptions(options));
         }
 
         setup.ReturnsAsync((new List<MCPServerConfig>(), new Dictionary<string, IList<McpClientTool>>()));
@@ -113,6 +114,7 @@ private void SetupMocksForToolEnumeration(List<MCPServerConfig> servers, Diction
             .Setup(x => x.EnumerateToolsFromServersAsync(
                 It.IsAny<string>(),
                 It.IsAny<string>(),
+                It.IsAny<TokenProviderCollection>(),
                 It.IsAny<ITurnContext>(),
                 It.IsAny<ToolOptions>()))
             .ReturnsAsync((servers, toolsByServer));
@@ -196,6 +198,7 @@ await _service.GetMcpToolDefinitionsAndResourcesAsync(
             x => x.EnumerateToolsFromServersAsync(
                 TestAgentInstanceId,
                 TestAuthToken,
+                It.IsAny<IMcpTokenProvider>(),
                 _mockTurnContext.Object,
                 It.Is<ToolOptions>(o => o.UserAgentConfiguration == Agent365AzureAIFoundrySdkUserAgentConfiguration.Instance)),
             Times.Once);

src/Tooling/Core/Microsoft.Agents.A365.Tooling.csproj

@@ -20,6 +20,7 @@
     <InternalsVisibleTo Include="Microsoft.Agents.A365.Tooling.Extensions.SemanticKernel" />
     <InternalsVisibleTo Include="Microsoft.Agents.A365.Tooling.Extensions.AgentFramework" />
     <InternalsVisibleTo Include="Microsoft.Agents.A365.Tooling.Extensions.AzureAIFoundry" />
+    <InternalsVisibleTo Include="Microsoft.Agents.A365.Tooling.Extensions.AzureAIFoundry.Tests" />
   </ItemGroup>
   <ItemGroup>
     <PackageReference Include="Microsoft.Extensions.Logging.Abstractions" />

src/Tooling/Core/Services/DevMcpTokenProvider.cs → src/Tooling/Core/Services/EnvMcpTokenProvider.cs

@@ -25,17 +25,17 @@ namespace Microsoft.Agents.A365.Tooling.Services
     /// </list>
     /// Hyphens in the server name are normalised to underscores before the lookup.
     /// </remarks>
-    internal sealed class DevMcpTokenProvider : IMcpTokenProvider
+    internal sealed class EnvMcpTokenProvider : IMcpTokenProvider
     {
         private readonly IConfiguration _configuration;
         private readonly ILogger _logger;
 
         /// <summary>
-        /// Initializes a new instance of <see cref="DevMcpTokenProvider"/>.
+        /// Initializes a new instance of <see cref="EnvMcpTokenProvider"/>.
         /// </summary>
         /// <param name="configuration">Application configuration (env vars, appsettings, etc.).</param>
         /// <param name="logger">Logger for diagnostic output.</param>
-        public DevMcpTokenProvider(IConfiguration configuration, ILogger logger)
+        public EnvMcpTokenProvider(IConfiguration configuration, ILogger logger)
         {
             _configuration = configuration;
             _logger = logger;
@@ -55,9 +55,13 @@ public Task<string> GetTokenAsync(MCPServerConfig server, CancellationToken canc
 
             if (string.IsNullOrWhiteSpace(token))
             {
-                throw new InvalidOperationException(
-                    $"No dev token found for MCP server '{server.mcpServerName}'. " +
-                    $"Set environment variable '{perServerKey}' or 'BEARER_TOKEN'.");
+                if (Utility.IsDevScenario(_configuration)) // Only log warnings in dev scenarios, to avoid noise in prod if env vars are not set
+                {
+                    _logger.LogWarning(
+                     $"No Environment token found for MCP server '{server.mcpServerName}'. " +
+                    $"Set environment variable '{perServerKey}' or 'BEARER_TOKEN'.");
+                }
+                return Task.FromResult(string.Empty);
             }
 
             // Warn when a V2 server (distinct audience) falls back to the shared BEARER_TOKEN.

src/Tooling/Core/Services/McpToolServerConfigurationService.cs

@@ -552,22 +552,39 @@ private async Task AttachPerAudienceTokensAsync(
             // Sequential acquisition avoids throttling the OBO endpoint.
             var tokenByScope = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
 
+            List<MCPServerConfig> failedToAcquireServers = new List<MCPServerConfig>();
             foreach (var server in servers)
             {
                 cancellationToken.ThrowIfCancellationRequested();
+                try
+                {
+                    var scope = Utils.Utility.ResolveTokenScopeForServer(server, _configuration);
+                    if (!tokenByScope.TryGetValue(scope, out var token))
+                    {
+                        token = await tokenProvider.GetTokenAsync(server, cancellationToken).ConfigureAwait(false);
+                        tokenByScope[scope] = token;
+                        _logger.LogDebug(
+                            "Acquired token for scope '{Scope}' (server '{ServerName}')",
+                            scope, server.mcpServerName);
+                    }
 
-                var scope = Utils.Utility.ResolveTokenScopeForServer(server, _configuration);
-                if (!tokenByScope.TryGetValue(scope, out var token))
+                    server.Headers ??= new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
+                    server.Headers[Constants.Headers.Authorization] = $"{Constants.Headers.BearerPrefix} {token}";
+                }
+                catch (Exception ex)
                 {
-                    token = await tokenProvider.GetTokenAsync(server, cancellationToken).ConfigureAwait(false);
-                    tokenByScope[scope] = token;
-                    _logger.LogDebug(
-                        "Acquired token for scope '{Scope}' (server '{ServerName}')",
-                        scope, server.mcpServerName);
+                    failedToAcquireServers.Add(server);
+                    _logger.LogError(ex, "Failed to acquire token for server '{ServerName}': {Message}", server.mcpServerName, ex.Message);
                 }
+            }
 
-                server.Headers ??= new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
-                server.Headers[Constants.Headers.Authorization] = $"{Constants.Headers.BearerPrefix} {token}";
+            if (failedToAcquireServers.Count > 0)
+            {
+                _logger.LogWarning("Failed to acquire tokens for {Count} MCP servers: {ServerNames}",
+                    failedToAcquireServers.Count, string.Join(", ", failedToAcquireServers.Select(s => s.mcpServerName)));
+                // remove servers we failed to acquire tokens for, since they'll likely fail authentication anyway
+                servers.RemoveAll(s => failedToAcquireServers.Contains(s));
+                failedToAcquireServers.Clear(); 
             }
         }
 

src/Tooling/Core/Services/TokenProviderCollection.cs

@@ -0,0 +1,66 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using Microsoft.Agents.A365.Tooling.Models;
+using Microsoft.Extensions.Logging;
+
+namespace Microsoft.Agents.A365.Tooling.Services
+{
+    internal class TokenProviderCollection : IMcpTokenProvider
+    {
+        readonly SortedDictionary<int, IMcpTokenProvider> _providers;
+        readonly ILogger _logger;
+
+        public TokenProviderCollection(
+            ILogger logger,
+            params IMcpTokenProvider[] providers)
+        {
+            _logger = logger;
+            _providers = new SortedDictionary<int, IMcpTokenProvider>();
+            for (int i = 0; i < providers.Length; i++)
+            {
+                _providers.Add(i, providers[i]);
+            }
+
+        }
+
+        public async Task<string> GetTokenAsync(MCPServerConfig server, CancellationToken cancellationToken = default)
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+
+            if ( _providers != null && _providers.Count == 0)
+            {
+                throw new InvalidOperationException("No token providers are configured.");
+            }
+
+            if (_providers!.Values.Any(p => p == null))
+            {
+                throw new InvalidOperationException("One or more token providers are null.");
+            }
+
+            List<Exception> exceptions = new List<Exception>();
+            // Try each provider in turn, then return the first successful token. If no providers can return a token, throw an exception.
+            foreach (var provider in _providers.Values)
+            {
+                try
+                {
+                    var token = await provider.GetTokenAsync(server, cancellationToken);
+                    if (!string.IsNullOrWhiteSpace(token))
+                    {
+                        return token;
+                    }
+                    else
+                    {
+                        _logger.LogDebug("Provider {ProviderName} returned an empty token.", provider.GetType().Name);
+                    }
+                }
+                catch(Exception ex)
+                {
+                    exceptions.Add(new Exception($"Provider {provider.GetType().Name} failed to obtain a token." , ex));
+                    _logger.LogDebug(ex, "Token provider {ProviderType} failed to obtain a token for server '{ServerName}'.", provider.GetType().Name, server.mcpServerName);
+                }
+            }
+            throw new AggregateException("No valid token could be obtained from any provider.", exceptions);
+        }
+    }
+}

src/Tooling/Extensions/AgentFramework/Services/McpToolRegistrationService.cs

@@ -3,25 +3,21 @@
 
 namespace Microsoft.Agents.A365.Tooling.Extensions.AgentFramework.Services;
 
-using System;
-using System.Collections.Generic;
-using System.Linq;
-using System.Threading;
-using System.Threading.Tasks;
 using Microsoft.Agents.A365.Runtime;
-using Microsoft.Agents.A365.Runtime.Authentication;
 using Microsoft.Agents.A365.Tooling.Models;
 using Microsoft.Agents.A365.Tooling.Services;
-using AgenticMcpTokenProvider = Microsoft.Agents.A365.Tooling.Services.AgenticMcpTokenProvider;
-using DevMcpTokenProvider = Microsoft.Agents.A365.Tooling.Services.DevMcpTokenProvider;
-using IMcpTokenProvider = Microsoft.Agents.A365.Tooling.Services.IMcpTokenProvider;
-using ToolingUtility = Microsoft.Agents.A365.Tooling.Utils.Utility;
 using Microsoft.Agents.AI;
 using Microsoft.Agents.Builder;
 using Microsoft.Agents.Builder.App.UserAuth;
 using Microsoft.Extensions.AI;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.Logging;
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using IMcpTokenProvider = Microsoft.Agents.A365.Tooling.Services.IMcpTokenProvider;
 
 /// <summary>
 /// Service for registering and validating MCP tool servers for Agent Framework scenarios.
@@ -64,10 +60,6 @@ public async Task<AIAgent> AddToolServersToAgent(
             throw new ArgumentNullException(nameof(chatClient));
         }
 
-        if (authToken is null && !ToolingUtility.IsDevScenario(_configuration))
-        {
-            authToken = await AgenticAuthenticationService.GetAgenticUserTokenAsync(userAuthorization, authHandlerName, turnContext, _configuration).ConfigureAwait(false);
-        }
         authToken ??= string.Empty;
 
         try
@@ -83,12 +75,19 @@ public async Task<AIAgent> AddToolServersToAgent(
                 UserAgentConfiguration = Agent365AgentFrameworkSdkUserAgentConfiguration.Instance
             };
 
-            // Use per-audience token provider so V2 servers receive audience-scoped tokens.
-            // In dev scenarios tokens come from environment variables; in production from OBO flow.
-            IMcpTokenProvider tokenProvider = ToolingUtility.IsDevScenario(_configuration)
-                ? new DevMcpTokenProvider(_configuration, _logger)
-                : new AgenticMcpTokenProvider(userAuthorization, authHandlerName, turnContext, _configuration, _logger);
-
+            IMcpTokenProvider tokenProvider;
+            if (userAuthorization is not null && authHandlerName is not null)
+            {
+                // Production V2-aware path: per-audience OBO tokens.
+                tokenProvider = new TokenProviderCollection(_logger,
+                        new EnvMcpTokenProvider(_configuration, _logger),
+                        new AgenticMcpTokenProvider(userAuthorization, authHandlerName, turnContext, _configuration, _logger));
+            }
+            else
+            {
+                tokenProvider = new TokenProviderCollection(_logger, new EnvMcpTokenProvider(_configuration, _logger));
+            }
+
             var (_, toolsByServer) = await _mcpServerConfigurationService.EnumerateToolsFromServersAsync(agentUserId, authToken, tokenProvider, turnContext, toolOptions).ConfigureAwait(false);
 
             // Add all MCP tools from all servers
@@ -126,20 +125,25 @@ public async Task<IList<AITool>> GetMcpToolsAsync(
     {
         try
         {
-            if (authToken is null && !ToolingUtility.IsDevScenario(_configuration))
-            {
-                authToken = await AgenticAuthenticationService.GetAgenticUserTokenAsync(userAuthorization, authHandlerName, turnContext, _configuration).ConfigureAwait(false);
-            }
             authToken ??= string.Empty;
 
             var toolOptions = new ToolOptions
             {
                 UserAgentConfiguration = Agent365AgentFrameworkSdkUserAgentConfiguration.Instance
             };
 
-            IMcpTokenProvider tokenProvider = ToolingUtility.IsDevScenario(_configuration)
-                ? new DevMcpTokenProvider(_configuration, _logger)
-                : new AgenticMcpTokenProvider(userAuthorization, authHandlerName, turnContext, _configuration, _logger);
+            IMcpTokenProvider tokenProvider;
+            if (userAuthorization is not null && authHandlerName is not null)
+            {
+                // Production V2-aware path: per-audience OBO tokens.
+                tokenProvider = new TokenProviderCollection(_logger,
+                        new EnvMcpTokenProvider(_configuration, _logger),
+                        new AgenticMcpTokenProvider(userAuthorization, authHandlerName, turnContext, _configuration, _logger));
+            }
+            else
+            {
+                tokenProvider = new TokenProviderCollection(_logger, new EnvMcpTokenProvider(_configuration, _logger));
+            }
 
             var (_, toolsByServer) = await _mcpServerConfigurationService.EnumerateToolsFromServersAsync(
                 agentUserId, authToken, tokenProvider, turnContext, toolOptions).ConfigureAwait(false);