feat: Dynamic authentication handling in MCPToolset from ADK Readonly… #2208

pandasanjay · 2025-07-28T08:49:05Z

…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:
- 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. - MCP Tools: No mechanism to pass user JWT token through ADK session.state to MCP server context #1402
- 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. - Implement outbound MCP authentication #805
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.

…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.

Copilot

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_callback in 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

Copilot · 2025-07-28T08:50:42Z

tests/unittests/tools/mcp_tool/test_mcp_toolset_env.py

+# Import after mocking to avoid MCP dependency issues
+from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset
+
+


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.

Suggested change

# Import after mocking to avoid MCP dependency issues

from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset

Copilot · 2025-07-28T08:50:42Z

tests/unittests/tools/mcp_tool/test_mcp_toolset_auth.py

+# Import after mocking to avoid MCP dependency issues
+from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset
+
+


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.

Suggested change

# Import after mocking to avoid MCP dependency issues

from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset

Copilot · 2025-07-28T08:50:42Z

tests/unittests/tools/mcp_tool/test_mcp_tool_auth.py

+# Import after mocking to avoid MCP dependency issues
+from google.adk.tools.mcp_tool.mcp_tool import MCPTool
+
+


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.

Suggested change

# Import after mocking to avoid MCP dependency issues

from google.adk.tools.mcp_tool.mcp_tool import MCPTool

Copilot · 2025-07-28T08:50:43Z

tests/integration/utils/test_runner.py

+    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


[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.

src/google/adk/tools/mcp_tool/mcp_toolset.py

pandasanjay · 2025-07-28T09:01:59Z

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.

ryanaiagent · 2025-12-04T23:56:51Z

Hi @pandasanjay , Thank you for your contribution through this pull request! I appreciate your patience with the delay. This PR has merge conflicts that require changes from your end. Could you please rebase your branch with the latest main branch to address these? Once this is complete, please let us know so we can proceed with the review.

Resolved merge conflicts in mcp_toolset.py and pyproject.toml. Merged main branch changes (McpToolset rename, retry_on_errors decorator, header_provider, from_config method) with dynamic MCP auth feature.

The method is already defined in BaseToolset and handles all cases.

pandasanjay · 2025-12-06T21:54:45Z

Hi @ryanaiagent, I’ve updated the code and would like you to review it.

I believe we now have a header_provider that solves the issue of passing SSE/Streamable HTTP headers. However, since my code was built before the introduction of the header_provider, it uses a different callback. I feel this is unnecessary, but I’ll remove it once you review the PR.

Additionally, this PR also includes another fix: passing an environment variable from the context state to the STDIO MCP.

ryanaiagent · 2025-12-30T21:18:16Z

/gemini review

gemini-code-assist

Code Review

This PR introduces a powerful feature for dynamic authentication by allowing auth details and environment variables to be passed from the ADK session state to MCP. The implementation in MCPToolset is well-structured, with new callbacks for extracting auth and environment data. The addition of comprehensive unit tests for the new auth logic is also a great improvement.

However, there are several critical issues that need addressing. The new integration tests and the new sample code are broken due to the refactoring that moved logic from MCPSessionManager to MCPToolset. They are calling non-existent methods and will fail. There's also a bug in MCPSessionManager.update_connection_params where it fails to clear the session cache. Additionally, there's a performance concern in MCPToolset.get_tools where a new session manager (and thus a new process) is created on every call with dynamic env vars. Finally, the new sample's README contains a hardcoded path.

I've provided detailed comments and suggestions for each of these points. Once these are addressed, this will be a solid contribution.

gemini-code-assist · 2025-12-30T21:20:02Z

contributing/samples/mcp_stdio_user_auth_passing_sample/main.py

+  mcp_session_manager = mcp_toolset._mcp_session_manager
+
+  # Extract environment variables using the callback (without connecting to MCP)
+  if mcp_session_manager._context_to_env_mapper_callback:
+    print("✓ Context-to-env mapper callback is configured")
+
+    # Simulate what happens during MCP session creation
+    env_vars = mcp_session_manager._extract_env_from_context(readonly_context)


The sample code attempts to access _mcp_session_manager._context_to_env_mapper_callback and call _mcp_session_manager._extract_env_from_context. These attributes and methods have been removed from MCPSessionManager and the logic moved to MCPToolset in this PR. As a result, the sample is broken and will not demonstrate the intended feature.

The sample should be updated to use the public APIs and reflect the new implementation. A better sample would perform an actual tool call to show the end-to-end flow, rather than relying on private members which can change.

gemini-code-assist · 2025-12-30T21:20:02Z

tests/integration/test_mcp_env_integration.py

+  async def test_session_manager_direct_env_injection(self, llm_backend):
+    """Test MCPSessionManager environment variable injection directly."""
+    connection_params = StdioServerParameters(
+        command='npx',
+        args=['-y', '@modelcontextprotocol/server-filesystem'],
+        env={'EXISTING_VAR': 'existing_value'},
+    )
+
+    def test_env_callback(state_dict: Dict[str, Any]) -> Dict[str, str]:
+      return {'NEW_VAR': 'new_value', 'API_KEY': 'secret123'}
+
+    session_manager = MCPSessionManager(
+        connection_params=connection_params,
+        get_env_from_context_fn=test_env_callback,
+    )


This test, and others in this file, appear to be broken after the refactoring. MCPSessionManager is initialized with get_env_from_context_fn, which is no longer a valid argument. Furthermore, the test attempts to call _extract_env_from_context and _inject_env_vars on the session_manager instance, but these methods do not exist on MCPSessionManager anymore; this logic has been moved to MCPToolset.

The entire test suite in this file needs to be updated to reflect the new architecture where MCPToolset is responsible for environment variable handling.

gemini-code-assist · 2025-12-30T21:20:02Z

src/google/adk/tools/mcp_tool/mcp_session_manager.py

+    # Clear existing sessions since connection params changed
+    # Sessions will be recreated on next request
+    # Note: We don't close sessions here to avoid blocking,
+    # they will be cleaned up when detected as disconnected


The docstring for update_connection_params states that it "invalidates existing sessions". However, the implementation only updates self._connection_params. The session cache self._sessions is not actually cleared, which means existing sessions with outdated connection parameters might be reused. This is a bug.

To fix this, you should clear the session cache.

Suggested change

# Clear existing sessions since connection params changed

# Sessions will be recreated on next request

# Note: We don't close sessions here to avoid blocking,

# they will be cleaned up when detected as disconnected

# Clear existing sessions since connection params changed

self._sessions.clear()

# Sessions will be recreated on next request

# Note: We don't close sessions here to avoid blocking,

# they will be cleaned up when detected as disconnected

gemini-code-assist · 2025-12-30T21:20:02Z

contributing/samples/mcp_stdio_user_auth_passing_sample/README.md

+cd /home/sanjay-dev/Workspace/adk-python
+python -m contributing.samples.stdio_mcp_user_auth_passing_sample.main


The command to run the sample includes a hardcoded, user-specific path (/home/sanjay-dev/Workspace/adk-python). This will not work for other users and will cause confusion. The instructions should be generic.

Suggested change

cd /home/sanjay-dev/Workspace/adk-python

python -m contributing.samples.stdio_mcp_user_auth_passing_sample.main

cd <path_to_adk-python_repo_root>

python -m contributing.samples.stdio_mcp_user_auth_passing_sample.main

gemini-code-assist · 2025-12-30T21:20:02Z

contributing/samples/mcp_stdio_user_auth_passing_sample/agent.py

+      instruction=f"""
+        You are an agent that calls an internal MCP server which requires a user token for internal API calls.
+        The user token is available in your session state and must be passed to the MCP process as an environment variable.
+        Test directory: {temp_dir}
+        """,


The agent's instruction string includes Test directory: {temp_dir}, but the StdioServerParameters for the filesystem server is configured with _allowed_path (on line 51). The temp_dir variable is created but not actually used by the MCP tool. This is misleading for a sample and could confuse users about how to configure the tool's working directory.

To make the sample clearer, I recommend using temp_dir as the path for StdioServerParameters to align the instruction with the implementation.

gemini-code-assist · 2025-12-30T21:20:02Z

src/google/adk/tools/mcp_tool/mcp_toolset.py

+    # If we have env vars to inject for STDIO, create a temporary session manager
+    if env_vars and isinstance(self._connection_params, StdioConnectionParams):
+      updated_connection_params = self._create_updated_connection_params(
+          env_vars
+      )
+      session_manager = MCPSessionManager(
+          connection_params=updated_connection_params,
+          errlog=self._errlog,
+      )


The get_tools method creates a new, temporary MCPSessionManager instance every time it's called with dynamic environment variables for an Stdio connection. This is inefficient because it will likely spawn a new MCP process on each call, instead of reusing an existing one. This could lead to significant performance overhead, especially if get_tools is called frequently.

Consider refactoring to avoid creating a new MCPSessionManager. One approach could be to enhance MCPSessionManager to manage sessions with different environment variables within the same manager instance, for example by including a hash of the environment variables in the session_key for stdio connections.

Copilot AI review requested due to automatic review settings July 28, 2025 08:49

Copilot AI reviewed Jul 28, 2025

View reviewed changes

pandasanjay mentioned this pull request Jul 28, 2025

Feat: Dynamic authentication handling in MCPToolset from Adk ReadonlyContext.state using Callback #1198

Closed

timof1308 mentioned this pull request Nov 23, 2025

feat(mcp): Add session state-based JWT token propagation for MCP tools #3675

Open

9 tasks

ryanaiagent self-assigned this Dec 1, 2025

ryanaiagent added mcp [Component] Issues about MCP support request clarification [Status] The maintainer need clarification or more information from the author labels Dec 4, 2025

pandasanjay added 2 commits December 6, 2025 15:38

chore(tools): merge main into feat-dynamic-mcp-auth-clean

bbf920e

Resolved merge conflicts in mcp_toolset.py and pyproject.toml. Merged main branch changes (McpToolset rename, retry_on_errors decorator, header_provider, from_config method) with dynamic MCP auth feature.

refactor(tools): remove redundant _is_tool_selected override

dd5ffa4

The method is already defined in BaseToolset and handles all cases.

gemini-code-assist bot reviewed Dec 30, 2025

View reviewed changes

		# Import after mocking to avoid MCP dependency issues
		from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset

		# Import after mocking to avoid MCP dependency issues
		from google.adk.tools.mcp_tool.mcp_tool import MCPTool

		cd /home/sanjay-dev/Workspace/adk-python
		python -m contributing.samples.stdio_mcp_user_auth_passing_sample.main

feat: Dynamic authentication handling in MCPToolset from ADK Readonly… #2208

Are you sure you want to change the base?

feat: Dynamic authentication handling in MCPToolset from ADK Readonly… #2208

Uh oh!

Conversation

pandasanjay commented Jul 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Copilot AI Jul 28, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 28, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 28, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jul 28, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

pandasanjay commented Jul 28, 2025

Uh oh!

ryanaiagent commented Dec 4, 2025

Uh oh!

pandasanjay commented Dec 6, 2025

Uh oh!

ryanaiagent commented Dec 30, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Dec 30, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Dec 30, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Dec 30, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Dec 30, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Dec 30, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Dec 30, 2025

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pandasanjay commented Jul 28, 2025 •

edited

Loading