Skip to content

MCPToolset: Add OAuth2 Client Credentials Flow with RFC 8414 Compliant Discovery #2061

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 10 commits into
base: main
Choose a base branch
from

Conversation

jeremyschulman
Copy link

@jeremyschulman jeremyschulman commented Jul 20, 2025

Add OAuth2 Client Credentials Flow with Automatic Discovery to ADK MCPToolset

The ADK now offers OAuth2 client credentials authentication that "just works" out of the box! 🚀

Summary

This PR adds enterprise-grade OAuth2 client credentials authentication to the ADK MCPToolset with automatic discovery capabilities. The implementation follows RFC 8414 for OAuth2 Authorization Server Metadata discovery and provides seamless authentication setup with minimal configuration.

No breaking changes - This PR is fully backwards compatible. Existing MCPToolset usage continues to work unchanged.

Key Features

🚀 Automatic OAuth Discovery by Default

  • MCPToolset now automatically enables OAuth discovery for HTTP-based connections
  • Zero configuration required for standard OAuth2 setups
  • Automatic base URL extraction from connection parameters

🔧 RFC 8414 Compliant Discovery

  • Two-stage OAuth discovery process:
    1. Query .well-known/oauth-protected-resource to find authorization server
    2. Query authorization server's .well-known/oauth-authorization-server for token endpoint
  • Fallback mechanisms for non-compliant servers

🔐 Complete OAuth2 Client Credentials Flow

  • Full token exchange implementation using authlib
  • Session-level authentication for MCP tool operations
  • Automatic token refresh and credential management

📝 Production-Ready Logging

  • Comprehensive DEBUG level tracing for development
  • Clean production logs with appropriate WARNING/ERROR levels
  • Full OAuth flow visibility when needed

Files Updated

🆕 New Files Created

adk-python/src/google/adk/auth/oauth2_discovery_util.py (~180 lines)

Purpose: OAuth2 discovery utilities following RFC 8414

Key Functions:

  • create_oauth_scheme_from_discovery() - Main discovery entry point
  • _query_oauth_protected_resource() - Query .well-known/oauth-protected-resource
  • _query_authorization_server_metadata() - Query authorization server metadata
  • _create_oauth2_scheme() - Create OAuth2 scheme from discovered endpoints

adk-python/src/google/adk/tools/mcp_tool/mcp_auth_discovery.py (~85 lines)

Purpose: Configuration class for OAuth discovery

Key Features:

  • MCPAuthDiscovery dataclass with base_url, timeout, enabled properties
  • Input validation and normalization
  • Clean API for discovery configuration

📝 Enhanced Existing Files

adk-python/src/google/adk/tools/mcp_tool/mcp_toolset.py (~100 lines added/modified)

Major Changes:

  • Default OAuth Discovery: Automatically creates MCPAuthDiscovery from connection params
  • _create_default_auth_discovery(): Extracts base URLs from HTTP connections
  • _perform_oauth_discovery(): Performs RFC 8414 two-stage discovery
  • Token Exchange: Exchanges OAuth2 credentials before session creation
  • Updated Constructor: New auth_discovery parameter with automatic defaults
  • Enhanced Documentation: Examples showing automatic vs explicit discovery

adk-python/src/google/adk/auth/credential_manager.py (~50 lines added/modified)

Key Enhancements:

  • OAuth2CredentialExchanger Registration: Registered in _exchanger_registry
  • Fallback Logic: Raw OAuth2 credential fallback for client credentials flow
  • Enhanced get_auth_credential(): 8-step credential processing workflow
  • Debug Logging: Comprehensive step-by-step flow tracing

adk-python/src/google/adk/auth/exchanger/oauth2_credential_exchanger.py (~80 lines added/modified)

Major Enhancements:

  • Client Credentials Support: _exchange_client_credentials() method
  • OAuth2Session Integration: Using authlib with client_secret_post method
  • Grant Type Detection: _get_grant_type() for flow type determination
  • Token Exchange: Full OAuth2 token exchange implementation
  • Error Handling: Comprehensive error handling and fallbacks

adk-python/src/google/adk/tools/mcp_tool/mcp_tool.py (~10 lines added/modified)

Minor Enhancements:

  • Debug Logging: Constructor logging for auth configuration
  • OAuth2 Header Support: Enhanced _get_headers() for OAuth2 tokens

Usage Examples

Automatic Discovery

# Automatic OAuth discovery - just works!
mcp_toolset = MCPToolset(
    connection_params=StreamableHTTPConnectionParams(
        url="http://localhost:9204/mcp/",
    ),
    auth_scheme=OAuth2(
        flows=OAuthFlows(
            clientCredentials=OAuthFlowClientCredentials(
                tokenUrl="",  # Empty - automatically discovered
                scopes={"api:read": "Read access"}
            )
        )
    ),
    auth_credential=oauth2_credential,
    # ✅ No auth_discovery needed - automatically enabled!
)

Override Options

# Custom discovery server
mcp_toolset = MCPToolset(
    connection_params=StreamableHTTPConnectionParams(...),
    auth_credential=oauth2_credential,
    auth_discovery=MCPAuthDiscovery(base_url='http://auth-server:9205')
)

# Disable discovery completely  
mcp_toolset = MCPToolset(
    connection_params=StreamableHTTPConnectionParams(...),
    auth_credential=oauth2_credential,
    auth_discovery=MCPAuthDiscovery(enabled=False)
)

Integration Flow

The implementation follows a clean integration pattern:

  1. MCPToolset → Creates default MCPAuthDiscovery from connection params
  2. oauth2_discovery_util → Discovers OAuth endpoints via RFC 8414
  3. credential_manager → Registers and uses OAuth2CredentialExchanger
  4. oauth2_credential_exchanger → Performs client credentials token exchange
  5. mcp_tool → Uses exchanged tokens for authenticated MCP requests

Default Behavior Logic

Auto-Enable OAuth Discovery:

  • StreamableHTTP: http://localhost:9204/mcp/http://localhost:9204 (discovery base)
  • SSE: http://server:8080/sse/http://server:8080 (discovery base)
  • Stdio: Disabled (local processes don't need OAuth)

Discovery Process:

  1. Extract base URL from connection parameters automatically
  2. Query .well-known/oauth-protected-resource for authorization server
  3. Query authorization server's .well-known/oauth-authorization-server for token endpoint
  4. Create OAuth2 scheme with discovered token endpoint
  5. Exchange client credentials for access token
  6. Use access token for authenticated MCP operations

Benefits

  1. Zero Configuration: OAuth discovery works out-of-the-box for HTTP connections
  2. Convention over Configuration: Sensible defaults with explicit override options
  3. Backwards Compatible: Existing explicit auth_discovery parameters still work
  4. Predictable: No complex fallback logic - clear default behavior
  5. Developer Friendly: Less boilerplate, fewer configuration mistakes
  6. Production Ready: Appropriate logging levels and error handling
  7. Standards Compliant: Follows RFC 8414 OAuth2 Authorization Server Metadata

Technical Details

OAuth2 Discovery Implementation

  • Two-stage discovery following RFC 8414 specification
  • Graceful fallbacks for servers with non-standard configurations
  • Timeout handling with configurable discovery timeouts
  • Error resilience with comprehensive exception handling

Authentication Flow

  • Client credentials grant type using authlib OAuth2Session
  • client_secret_post authentication method (credentials in form body)
  • Session pooling with authentication headers
  • Automatic token refresh through credential manager

Logging Strategy

  • DEBUG: OAuth discovery flow details, token exchange progress, tool creation steps
  • WARNING: Configuration issues, missing auth schemes, fallback scenarios
  • ERROR: Authentication failures, missing required parameters, exchange errors

Code Statistics

  • Total Lines Added/Modified: ~505 lines across 6 ADK files
  • New Files: 2 (~265 lines)
  • Enhanced Files: 4 (~240 lines modified)
  • Test Coverage: OAuth discovery and client credentials flow

Test Coverage

🧪 Comprehensive Test Suite Added

The OAuth2 enhancement includes comprehensive test coverage across all new functionality:

New Test Files Created

  1. test_mcp_auth_discovery.py (~140 lines)

    • Basic initialization and validation
    • URL normalization (trailing slashes, complex paths)
    • Input validation (empty URLs, negative timeouts)
    • Default behavior and property testing
    • Dataclass equality and string representation
  2. test_mcp_toolset_oauth_discovery.py (~380 lines)

    • Default OAuth discovery creation for different connection types
    • Explicit auth_discovery parameter override behavior
    • OAuth discovery process with empty tokenUrl scenarios
    • Discovery failure and exception handling
    • URL parsing with complex paths and query parameters
    • Token exchange integration before session creation
    • Discovery only attempted once per toolset instance
  3. test_credential_manager_oauth2_integration.py (~140 lines)

    • OAuth2CredentialExchanger registration verification
    • Complete OAuth2 credential exchange flow testing
    • Raw OAuth2 credential fallback logic verification

Enhanced Existing Test Coverage

Existing comprehensive tests already cover:

  • test_oauth2_discovery_util.py - OAuth2 discovery utilities (RFC 8414)
  • test_oauth2_credential_exchanger.py - Client credentials flow implementation
  • test_credential_manager.py - Credential management workflows

🎯 Test Coverage Areas

  • OAuth Discovery Configuration - MCPAuthDiscovery validation and behavior
  • Automatic Discovery - Default behavior for different connection types
  • RFC 8414 Compliance - Two-stage discovery process
  • Client Credentials Flow - Token exchange using authlib
  • Error Handling - Discovery failures, network errors, invalid responses
  • URL Parsing - Base URL extraction from complex connection parameters
  • Session Authentication - Token exchange before MCP session creation
  • Credential Management - OAuth2CredentialExchanger registration and usage
  • Type Safety - Proper type annotations and validation
  • Edge Cases - Empty tokens, missing configurations, disabled discovery

📊 Test Statistics

  • Total Test Lines: ~660 lines of new test code
  • Test Methods: 25+ individual test methods
  • Coverage Areas: 8 major functional areas
  • Mock Integration: Comprehensive mocking of HTTP requests and OAuth sessions
  • Async Testing: Full async/await test coverage for discovery and exchange flows

Dependencies

  • authlib: Already in ADK dependencies for OAuth2Session support
  • httpx: Already in ADK dependencies for HTTP requests
  • fastapi.openapi.models: Already in ADK dependencies for OAuth2 schema types

Sample Implementation

🎯 Comprehensive OAuth2 Client Credentials Sample

A complete sample implementation has been created at:
adk-python/contributing/samples/mcp_oauth2_client_credentials_agent/

Sample Features:

  • 5 Different Scenarios: From automatic discovery to custom configurations
  • Mock OAuth2 Server: Complete RFC 8414 compliant test server
  • Production Examples: Real-world usage patterns and configurations
  • Comprehensive Documentation: Step-by-step setup and usage guide

Sample Contents:

  • agent.py - Five different OAuth2 agent configurations
  • mock_oauth_server.py - Complete OAuth2 test server implementation
  • README.md - Comprehensive documentation and usage guide
  • __init__.py - Standard Python package structure

Demonstrated Scenarios:

  1. Automatic Discovery - Zero configuration OAuth2 setup
  2. Custom Discovery - Override discovery server and settings
  3. Minimal Configuration - Let discovery create complete scheme
  4. Manual Configuration - Traditional OAuth2 setup without discovery
  5. Multi-server Setup - Multiple MCP servers with different auth configs

Mock OAuth2 Server Features:

  • RFC 8414 discovery endpoints (.well-known/oauth-protected-resource, .well-known/oauth-authorization-server)
  • Client credentials grant type support
  • Token validation endpoint for debugging
  • Multiple demo client configurations
  • Comprehensive logging for troubleshooting

Copy link

google-cla bot commented Jul 20, 2025

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@jeremyschulman jeremyschulman changed the title WIP -- ADK MCPToolset: Add OAuth2 Client Credentials Flow with Automatic Discovery WIP -- ADK MCPToolset: Add OAuth2 Client Credentials Flow with RFC 8414 Compliant Discovery Jul 20, 2025
@jeremyschulman jeremyschulman changed the title WIP -- ADK MCPToolset: Add OAuth2 Client Credentials Flow with RFC 8414 Compliant Discovery WIP -- MCPToolset: Add OAuth2 Client Credentials Flow with RFC 8414 Compliant Discovery Jul 20, 2025
@jeremyschulman jeremyschulman changed the title WIP -- MCPToolset: Add OAuth2 Client Credentials Flow with RFC 8414 Compliant Discovery MCPToolset: Add OAuth2 Client Credentials Flow with RFC 8414 Compliant Discovery Jul 20, 2025
@jeremyschulman jeremyschulman changed the title MCPToolset: Add OAuth2 Client Credentials Flow with RFC 8414 Compliant Discovery WIP: MCPToolset: Add OAuth2 Client Credentials Flow with RFC 8414 Compliant Discovery Jul 20, 2025
@jeremyschulman jeremyschulman changed the title WIP: MCPToolset: Add OAuth2 Client Credentials Flow with RFC 8414 Compliant Discovery MCPToolset: Add OAuth2 Client Credentials Flow with RFC 8414 Compliant Discovery Jul 20, 2025
@@ -307,6 +307,28 @@
* Added unit test coverage for local_eval_sets_manager.py ([174afb3](https://github.com/google/adk-python/commit/174afb3975bdc7e5f10c26f3eebb17d2efa0dd59))
* Extract common options for `adk web` and `adk api_server` ([01965bd](https://github.com/google/adk-python/commit/01965bdd74a9dbdb0ce91a924db8dee5961478b8))

## [Unreleased]
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

CHANGELOG is only updated up on release. please put those information in the message.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

will do!

@seanzhou1023
Copy link
Collaborator

seanzhou1023 commented Jul 22, 2025

Thank you so much for the amazing PR @jeremyschulman and I do think this feature is useful !!

However such PR is too big to review, would you please kindly split it into multiple smaller PRs (group those logically highly related changes into one PR) , thank you :)

@jeremyschulman
Copy link
Author

@seanzhou1023 - Could you please provide guidance on how you would like this PR split into parts? From a Developer perspective, the feature enhancement is all part-and-parcel together. Happy to split it up, and I value your help.

@georgerooney
Copy link

Hey @seanzhou1023 I'm working with Jeremy on this and wanted to propose this and get input for how to split.

  1. Auth Changes + associated tests
  2. MCP Server discovery + associated tests
  3. Contributor samples.

I know we could put 3 inside 1 and 2 partially but we're trying to be mindful on sizing of the PRs.

@seanzhou1023
Copy link
Collaborator

Hey @seanzhou1023 I'm working with Jeremy on this and wanted to propose this and get input for how to split.

  1. Auth Changes + associated tests
  2. MCP Server discovery + associated tests
  3. Contributor samples.

I know we could put 3 inside 1 and 2 partially but we're trying to be mindful on sizing of the PRs.

Never mind, I'm in the middle of review, it just takes longer to review big PRs.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants