Refactor MCP token provider selection and fallback logic#251
Open
MattB-msft wants to merge 3 commits into
Open
Refactor MCP token provider selection and fallback logic#251MattB-msft wants to merge 3 commits into
MattB-msft wants to merge 3 commits into
Conversation
Rename DevMcpTokenProvider to EnvMcpTokenProvider and update all usages. Introduce TokenProviderCollection to allow chaining multiple token providers, enabling fallback from environment-based to OBO-based authentication. Update token acquisition logic to handle failures more robustly by removing servers that fail token acquisition and logging warnings. Update tests and remove obsolete aliases. These changes improve flexibility, reliability, and diagnostics for MCP token handling.
Dependency Review✅ No vulnerabilities or license issues or OpenSSF Scorecard issues found.Scanned FilesNone |
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Pull request overview
This PR refactors MCP token acquisition to support chaining/fallback across multiple IMcpTokenProvider implementations (environment-based first, then OBO/agentic), and improves robustness by removing servers that fail per-audience token attachment.
Changes:
- Introduces
TokenProviderCollectionto try multiple token providers in order until a token is obtained. - Renames
DevMcpTokenProvidertoEnvMcpTokenProviderand updates extension registration services to use the chained provider approach. - Updates per-audience token attachment to log acquisition failures and remove servers that cannot be authenticated.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 10 comments.
Show a summary per file
| File | Description |
|---|---|
| src/Tooling/Extensions/SemanticKernel/Services/McpToolRegistrationService.cs | Switches to chained token providers for MCP tool enumeration. |
| src/Tooling/Extensions/AzureAIFoundry/Services/McpToolRegistrationService.cs | Uses chained token providers; adds dev/prod selection logic (with a nullability concern noted). |
| src/Tooling/Extensions/AgentFramework/Services/McpToolRegistrationService.cs | Switches to chained token providers for MCP tool enumeration. |
| src/Tooling/Core/Services/TokenProviderCollection.cs | Adds new chaining token provider implementation (several issues noted). |
| src/Tooling/Core/Services/McpToolServerConfigurationService.cs | Adds failure handling during token attachment by removing failing servers (cancellation and perf concerns noted). |
| src/Tooling/Core/Services/EnvMcpTokenProvider.cs | Renames the dev env-var provider class to EnvMcpTokenProvider. |
| src/Tests/Microsoft.Agents.A365.Tooling.Core.Tests/DevMcpTokenProviderTests.cs | Updates tests to target EnvMcpTokenProvider (naming consistency noted). |
Comment on lines
+1
to
+5
| using Microsoft.Agents.A365.Tooling.Models; | ||
| using System; | ||
| using System.Collections.Generic; | ||
| using System.Linq; | ||
| using System.Text; |
Comment on lines
+13
to
+20
| public TokenProviderCollection(params IMcpTokenProvider[] providers) | ||
| { | ||
| _providers = new SortedDictionary<int, IMcpTokenProvider>(); | ||
| for (int i = 0; i < providers.Length; i++) | ||
| { | ||
| _providers.Add(i, providers[i]); | ||
| } | ||
| } |
Comment on lines
176
to
+181
| if (ToolingUtility.IsDevScenario(_configuration)) | ||
| { | ||
| IMcpTokenProvider tokenProvider = new DevMcpTokenProvider(_configuration, _logger); | ||
| IMcpTokenProvider tokenProvider = new TokenProviderCollection( | ||
| new EnvMcpTokenProvider(_configuration, _logger), | ||
| new AgenticMcpTokenProvider(userAuthorization!, authHandlerName!, turnContext, _configuration, _logger)); | ||
|
|
Comment on lines
172
to
176
| // Select token provider: | ||
| // Dev scenario → DevMcpTokenProvider reads per-server env vars (no OBO flow needed). | ||
| // Production → AgenticMcpTokenProvider performs OBO when auth objects are supplied; | ||
| // falls back to the V1 shared-token path when they are absent. | ||
| if (ToolingUtility.IsDevScenario(_configuration)) |
| // Sequential acquisition avoids throttling the OBO endpoint. | ||
| var tokenByScope = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase); | ||
|
|
||
| List<MCPServerConfig> failedToAquireServers = new List<MCPServerConfig>(); |
Comment on lines
558
to
578
| 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); | ||
| failedToAquireServers.Add(server); | ||
| _logger.LogError(ex, "Failed to acquire token for server '{ServerName}': {Message}", server.mcpServerName, ex.Message); | ||
| } |
Comment on lines
+581
to
+587
| if (failedToAquireServers.Count > 0) | ||
| { | ||
| _logger.LogWarning("Failed to acquire tokens for {Count} MCP servers: {ServerNames}", | ||
| failedToAquireServers.Count, string.Join(", ", failedToAquireServers.Select(s => s.mcpServerName))); | ||
| // remove servers we failed to acquire tokens for, since they'll likely fail authentication anyway | ||
| servers.RemoveAll(s => failedToAquireServers.Contains(s)); | ||
| failedToAquireServers.Clear(); |
Comment on lines
+16
to
24
| /// Tests for <see cref="EnvMcpTokenProvider"/> and <see cref="Utility.IsDevScenario"/>. | ||
| /// </summary> | ||
| public class DevMcpTokenProviderTests | ||
| { | ||
| 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>()); |
biswapm
previously approved these changes
May 14, 2026
Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Comment on lines
+39
to
+42
| catch(Exception ex) | ||
| { | ||
| exceptions.Add(new Exception($"Provider {provider.GetType().Name} failed to obtain a token." , ex)); | ||
| } |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 7 changed files in this pull request and generated 11 comments.
Comments suppressed due to low confidence (2)
src/Tooling/Extensions/AgentFramework/Services/McpToolRegistrationService.cs:141
- Same issue as above: this TokenProviderCollection ordering will typically trigger an exception path on every token acquisition in production when env vars aren't set. Consider choosing provider order based on IsDevScenario, or avoiding exception-driven fallback.
IMcpTokenProvider tokenProvider = new TokenProviderCollection(
new EnvMcpTokenProvider(_configuration, _logger),
new AgenticMcpTokenProvider(userAuthorization, authHandlerName, turnContext, _configuration, _logger));
src/Tooling/Core/Services/McpToolServerConfigurationService.cs:587
- servers.RemoveAll(s => failedToAquireServers.Contains(s)) is O(n^2) due to List.Contains inside RemoveAll. Consider using a HashSet (or removing by index) to avoid quadratic behavior if the server list grows.
_logger.LogWarning("Failed to acquire tokens for {Count} MCP servers: {ServerNames}",
failedToAquireServers.Count, string.Join(", ", failedToAquireServers.Select(s => s.mcpServerName)));
// remove servers we failed to acquire tokens for, since they'll likely fail authentication anyway
servers.RemoveAll(s => failedToAquireServers.Contains(s));
failedToAquireServers.Clear();
Comment on lines
+1
to
+2
| // Copyright (c) Microsoft Corporation. | ||
| // Licensed under the MIT License. |
Comment on lines
+7
to
+8
| using System.Linq; | ||
| using System.Text; |
Comment on lines
+22
to
+42
| } | ||
| } | ||
|
|
||
| public async Task<string> GetTokenAsync(MCPServerConfig server, CancellationToken cancellationToken = default) | ||
| { | ||
| 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; | ||
| } | ||
| } | ||
| catch(Exception ex) | ||
| { | ||
| exceptions.Add(new Exception($"Provider {provider.GetType().Name} failed to obtain a token." , ex)); | ||
| } |
Comment on lines
73
to
+78
| // 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 = | ||
| new TokenProviderCollection( | ||
| new EnvMcpTokenProvider(_configuration, _logger), | ||
| new AgenticMcpTokenProvider(userAuthorization, authHandlerName, turnContext, _configuration, _logger)); |
Comment on lines
84
to
89
| // 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 = new TokenProviderCollection( | ||
| new EnvMcpTokenProvider(_configuration, _logger), | ||
| new AgenticMcpTokenProvider(userAuthorization, authHandlerName, turnContext, _configuration, _logger)); | ||
|
|
Comment on lines
176
to
+181
| if (ToolingUtility.IsDevScenario(_configuration)) | ||
| { | ||
| IMcpTokenProvider tokenProvider = new DevMcpTokenProvider(_configuration, _logger); | ||
| IMcpTokenProvider tokenProvider = new TokenProviderCollection( | ||
| new EnvMcpTokenProvider(_configuration, _logger), | ||
| new AgenticMcpTokenProvider(userAuthorization!, authHandlerName!, turnContext, _configuration, _logger)); | ||
|
|
Comment on lines
184
to
191
| else if (userAuthorization is not null && authHandlerName is not null) | ||
| { | ||
| // Production V2-aware path: per-audience OBO tokens. | ||
| IMcpTokenProvider tokenProvider = new AgenticMcpTokenProvider( | ||
| userAuthorization, authHandlerName, turnContext, _configuration, _logger); | ||
| IMcpTokenProvider tokenProvider = new TokenProviderCollection( | ||
| new EnvMcpTokenProvider(_configuration, _logger), | ||
| new AgenticMcpTokenProvider(userAuthorization, authHandlerName, turnContext, _configuration, _logger)); | ||
|
|
||
| (servers, toolsByServer) = await _mcpServerConfigurationService.EnumerateToolsFromServersAsync(agentInstanceId, authToken, tokenProvider, turnContext, toolOptions).ConfigureAwait(false); |
Comment on lines
558
to
578
| 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); | ||
| failedToAquireServers.Add(server); | ||
| _logger.LogError(ex, "Failed to acquire token for server '{ServerName}': {Message}", server.mcpServerName, ex.Message); | ||
| } |
Comment on lines
+555
to
556
| List<MCPServerConfig> failedToAquireServers = new List<MCPServerConfig>(); | ||
| foreach (var server in servers) |
| /// Tests for <see cref="DevMcpTokenProvider"/> and <see cref="Utility.IsDevScenario"/>. | ||
| /// Tests for <see cref="EnvMcpTokenProvider"/> and <see cref="Utility.IsDevScenario"/>. | ||
| /// </summary> | ||
| public class DevMcpTokenProviderTests |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Rename DevMcpTokenProvider to EnvMcpTokenProvider and update all usages. Introduce TokenProviderCollection to allow chaining multiple token providers, enabling fallback from environment-based to OBO-based authentication. Update token acquisition logic to handle failures more robustly by removing servers that fail token acquisition and logging warnings. Update tests and remove obsolete aliases. These changes improve flexibility, reliability, and diagnostics for MCP token handling.
This pull request introduces a new
TokenProviderCollectionto support chaining multiple token providers, improves error handling during token acquisition, and standardizes the naming of the environment-based token provider. The changes enhance flexibility and robustness in how tokens are acquired for MCP servers, especially across different environments (development and production).Token Provider Refactoring and Chaining:
TokenProviderCollectionclass, which allows chaining multipleIMcpTokenProviderimplementations and tries each in order until a valid token is obtained. This enables seamless fallback between environment-based and agentic token providers. (TokenProviderCollection.cssrc/Tooling/Core/Services/TokenProviderCollection.csR1-R44)DevMcpTokenProviderto use the newEnvMcpTokenProvider, and replaced direct instantiations with the newTokenProviderCollectionto allow for provider chaining in all registration services. [1] [2] [3] [4] [5] [6]Error Handling Improvements:
McpToolServerConfigurationService.AttachPerAudienceTokensAsyncby logging token acquisition failures per server, aggregating failed servers, and removing them from the list to prevent downstream authentication errors. [1] [2]Code Cleanup and Renaming:
DevMcpTokenProvidertoEnvMcpTokenProviderthroughout the codebase for clarity and consistency, and updated related references and comments. [1] [2]usingstatements for the old provider aliases in all affected files. [1] [2] [3]These changes collectively make the token acquisition process more robust, maintainable, and adaptable to different deployment scenarios.
OTEL telemetry indicating continue after tool lookup failure :
