feat: Dynamic authentication handling in MCPToolset from ADK Readonly… #2208
Conversation
…Context.state
This PR implements dynamic authentication handling in MCPToolset by passing
authentication information (such as user-specific tokens) directly from the
ADK session.state. With this approach, user tokens can be injected and sent
from the ADK Context state to MCP, supporting user-specific authentication flows.
Key Improvements:
1. Dynamic Token Management:
- MCPToolset now reads authentication information from the ADK ReadonlyContext.state
at runtime, enabling dynamic, user-specific token management.
- STDIO Transport: Authentication information is passed to MCP via environment variables.
- Other Transports (SSE and HTTP Stream): MCPToolset now receives authentication
details from the ADK ReadonlyContext.state, making them available for further
transmission, enabling a unified authentication mechanism across all transports.
2. Improved Unit Tests: Enhancements for better readability and maintainability.
3. New Sample: Demonstrates passing user tokens to MCP via environment variables,
showcasing the dynamic flow from session.state.
4. Flexible & Secure Integration: Allows for more flexible and secure integration
scenarios, supporting distinct authentication credentials per session or user.
There was a problem hiding this comment.
Pull Request Overview
This PR implements dynamic authentication handling in MCPToolset by passing authentication information directly from the ADK ReadonlyContext.state at runtime, enabling user-specific authentication flows and environment variable injection for MCP servers.
Key changes:
- Added dynamic token management through environment variable injection for STDIO transport and context-based authentication extraction for all transports
- Enhanced unit test coverage with comprehensive test suites for authentication and environment variable functionality
- Updated MCP dependency version to support new features
Reviewed Changes
Copilot reviewed 11 out of 12 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
src/google/adk/tools/mcp_tool/mcp_toolset.py |
Added auth extraction from context, environment variable injection, and callback support for dynamic authentication |
src/google/adk/tools/mcp_tool/mcp_session_manager.py |
Added connection parameter update functionality to support dynamic environment variable injection |
tests/unittests/tools/mcp_tool/test_mcp_toolset_env.py |
Comprehensive unit tests for environment variable extraction and injection functionality |
tests/unittests/tools/mcp_tool/test_mcp_toolset_auth.py |
Extensive unit tests for authentication extraction with error handling and validation |
tests/unittests/tools/mcp_tool/test_mcp_tool_auth.py |
Unit tests for MCPTool authentication functionality |
tests/integration/utils/test_runner.py |
Updated TestRunner with async support and session management improvements |
tests/integration/test_mcp_env_integration.py |
Integration tests demonstrating end-to-end environment variable and authentication flows |
pyproject.toml |
Updated MCP dependency version from 1.8.0 to 1.9.4 |
contributing/samples/mcp_stdio_user_auth_passing_sample/ |
Sample implementation showing user token passing through environment variables |
Comments suppressed due to low confidence (2)
tests/integration/test_mcp_env_integration.py:63
- Accessing private attributes like
_context_to_env_mapper_callbackin tests creates tight coupling and makes the code brittle to internal changes. Consider adding a public method or property to check if the callback is configured.
""",
pyproject.toml:41
- MCP version 1.9.4 may not exist. Please verify that this version is available before updating the dependency requirement.
"mcp>=1.9.4;python_version>='3.10'", # For MCP Toolset
| # Import after mocking to avoid MCP dependency issues | ||
| from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset | ||
|
|
||
|
|
There was a problem hiding this comment.
Import should be placed at the top of the file with other imports, not after mocking fixtures. Consider restructuring to avoid circular import issues while maintaining clean import organization.
| # Import after mocking to avoid MCP dependency issues | |
| from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset |
Copilot uses AI. Check for mistakes.
| # Import after mocking to avoid MCP dependency issues | ||
| from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset | ||
|
|
||
|
|
There was a problem hiding this comment.
Import should be placed at the top of the file with other imports, not after mocking fixtures. Consider restructuring to avoid circular import issues while maintaining clean import organization.
| # Import after mocking to avoid MCP dependency issues | |
| from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset |
Copilot uses AI. Check for mistakes.
| # Import after mocking to avoid MCP dependency issues | ||
| from google.adk.tools.mcp_tool.mcp_tool import MCPTool | ||
|
|
||
|
|
There was a problem hiding this comment.
Import should be placed at the top of the file with other imports, not after mocking fixtures. Consider restructuring to avoid circular import issues while maintaining clean import organization.
| # Import after mocking to avoid MCP dependency issues | |
| from google.adk.tools.mcp_tool.mcp_tool import MCPTool |
Copilot uses AI. Check for mistakes.
| self.current_session_id = None | ||
| self._session_initialized = False | ||
|
|
||
| def new_session(self, session_id: Optional[str] = None) -> None: | ||
| self.current_session_id = self.session_service.create_session( | ||
| async def _ensure_session(self) -> str: | ||
| """Ensure a session is created and return the session ID.""" | ||
| if not self._session_initialized: | ||
| session = await self.session_service.create_session( | ||
| app_name=self.app_name, user_id=self.user_id | ||
| ) | ||
| self.current_session_id = session.id | ||
| self._session_initialized = True | ||
| return self.current_session_id | ||
|
|
||
| async def new_session_async(self, session_id: Optional[str] = None) -> None: | ||
| session = await self.session_service.create_session( | ||
| app_name=self.app_name, user_id=self.user_id, session_id=session_id | ||
| ).id | ||
| ) | ||
| self.current_session_id = session.id | ||
| self._session_initialized = True |
There was a problem hiding this comment.
[nitpick] The initialization of current_session_id to None followed by session creation logic adds complexity. Consider initializing the session in __init__ or using a property to encapsulate the session creation logic more cleanly.
Copilot uses AI. Check for mistakes.
| from mcp import StdioServerParameters | ||
|
|
There was a problem hiding this comment.
Import statements should be placed at the top of the file, not inside methods. This import is used only once and could be moved to the module level imports.
| from mcp import StdioServerParameters |
Copilot uses AI. Check for mistakes.
|
Hi @hangfei, I've made the requested changes. Apologies for creating a new branch and pull request—I ran into some issues with the commit history on the previous one. This pull request contains the latest code along with all the fixes. |
…Context.state
This PR implements dynamic authentication handling in MCPToolset by passing authentication information (such as user-specific tokens) directly from the ADK session.state. With this approach, user tokens can be injected and sent from the ADK Context state to MCP, supporting user-specific authentication flows.
Key Improvements:
Dynamic Token Management:
Improved Unit Tests: Enhancements for better readability and maintainability.
New Sample: Demonstrates passing user tokens to MCP via environment variables, showcasing the dynamic flow from session.state.
Flexible & Secure Integration: Allows for more flexible and secure integration scenarios, supporting distinct authentication credentials per session or user.