Refactor MCP token provider selection and fallback logic#251
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 |
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). |
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>
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();
Corrected class name in tests to match the class being tested. Modified error handling in `EnvMcpTokenProvider` to log warnings instead of throwing exceptions in dev scenarios. Fixed typos in `McpToolServerConfigurationService`. Added logging to `TokenProviderCollection` and ensured token providers are configured and not null. Refactored token provider selection logic in `McpToolRegistrationService` to differentiate between production and development scenarios. Reorganized imports and updated token handling in `SemanticKernel` service.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 7 changed files in this pull request and generated 7 comments.
Comments suppressed due to low confidence (1)
src/Tooling/Core/Services/EnvMcpTokenProvider.cs:60
EnvMcpTokenProvider.GetTokenAsyncnow returns an empty string when no env token is present. This changes the provider contract (and will break existing tests/consumers that expect an exception for a missing token). Consider keeping the structured warning (in dev) but throwingInvalidOperationExceptionwhen no token is available, so the failure is explicit and TokenProviderCollection can still fall back to the next provider via its exception handling.
Rename DevMcpTokenProvider to EnvMcpTokenProvider Renamed `DevMcpTokenProvider` to `EnvMcpTokenProvider` across multiple files to reflect the use of environment variables for token provision in development scenarios. Updated comments in `EnvMcpTokenProviderTests.cs` to match the new naming. Added debug logging in `TokenProviderCollection.cs` to track when token providers return empty tokens or fail to obtain tokens, aiding in diagnostics. Removed unnecessary `using` directives in `McpToolRegistrationService.cs` and updated comments to reflect the renaming. Updated `CHANGELOG.md` to ensure documentation consistency with the code changes. These updates enhance code clarity, maintainability, and debugging capabilities.
Removed redundant test cases from EnvMcpTokenProviderTests that checked for exceptions with missing or whitespace-only token environment variables. Added a new mock field for TokenProviderCollection in McpToolRegistrationServiceTests and updated method signatures to include it. Updated the project file to grant internal access to a new test assembly for better testing coverage.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 10 out of 10 changed files in this pull request and generated 6 comments.
Comments suppressed due to low confidence (1)
src/Tooling/Core/Services/EnvMcpTokenProvider.cs:63
- Using string interpolation in ILogger.LogWarning prevents structured logging and adds avoidable allocations. Use a message template with parameters instead (also aligns with the rest of this file’s logging style).
Replaced `TokenProviderCollection` with `IMcpTokenProvider` in `McpToolRegistrationServiceTests.cs`. Removed `_mockTokenProviderCollection` field and its initialization. Updated `EnumerateToolsFromServersAsync` call to use `IMcpTokenProvider`, aligning with changes in the underlying service interface.
|
Also add test coverage for the following:
|
Restored logic to obtain authToken using AgenticAuthenticationService.GetAgenticUserTokenAsync when not provided but userAuthorization and authHandlerName are available. This supports both V1 and V2 authentication and provides fallback for non-Agent SDK contexts. Also, reorder using directives for consistency.
I am doing practical tests with V1 / v2 and blended scenarios as they don't make sense to mock full given that the functional code that has changed is token acquire and not server URL information. The scopes are either read from the config (V2) or they use the default scope / env provided scope. I will add missing tests for a few of the missing use cases. |
Refactored the selection of IMcpTokenProvider to use a single conditional (ternary) expression instead of if-else statements. This change reduces code duplication and improves readability while preserving the original logic for both development and production scenarios.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 10 out of 10 changed files in this pull request and generated 3 comments.
Comments suppressed due to low confidence (2)
src/Tooling/Core/Services/EnvMcpTokenProvider.cs:62
- This
LogWarninguses string interpolation, which defeats structured logging and allocates the message even when the warning level is disabled. Use a message template with arguments (consistent with the other log calls in this file).
src/Tests/Microsoft.Agents.A365.Tooling.Core.Tests/EnvMcpTokenProviderTests.cs:94 EnvMcpTokenProvidernow returnsstring.Emptywhen neitherBEARER_TOKEN_<SERVERNAME>norBEARER_TOKENis set (instead of throwing), but there’s no unit test asserting this new behavior. Adding a test prevents future refactors from accidentally reintroducing an exception or returning non-empty whitespace.
Simplified IMcpTokenProvider setup in McpToolRegistrationService by replacing the if/else block with a conditional expression. This reduces code duplication and clarifies provider selection based on userAuthorization and authHandlerName presence.
| 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; |
| cancellationToken.ThrowIfCancellationRequested(); | ||
| try | ||
| { | ||
| var scope = Utils.Utility.ResolveTokenScopeForServer(server, _configuration); | ||
| if (!tokenByScope.TryGetValue(scope, out var token)) |
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 :
