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

contributing/samples/mcp_stdio_user_auth_passing_sample/README.md

@@ -0,0 +1,80 @@
+# Sample: Passing User Token from Agent State to MCP via ContextToEnvMapperCallback
+
+This sample demonstrates how to use the `context_to_env_mapper_callback` feature in ADK to pass a user token from the agent's session state to an MCP process (using stdio transport). This is useful when your MCP server (built by your organization) requires the same user token for internal API calls.
+
+## How it works
+- The agent is initialized with a `MCPToolset` using `StdioServerParameters`.
+- The `context_to_env_mapper_callback` is set to a function that extracts the `user_token` from the agent's state and maps it to the `USER_TOKEN` environment variable.
+- When the agent calls the MCP, the token is injected into the MCP process environment, allowing the MCP to use it for internal authentication.
+
+## Directory Structure
+```
+contributing/samples/stdio_mcp_user_auth_passing_sample/
+├── agent.py       # Basic agent setup
+├── main.py        # Complete runnable example
+└── README.md
+```
+
+## How to Run
+
+### Option 1: Run the complete example
+```bash
+cd /home/sanjay-dev/Workspace/adk-python
+python -m contributing.samples.stdio_mcp_user_auth_passing_sample.main
+```
+
+### Option 2: Use the agent in your own code
+```python
+from contributing.samples.stdio_mcp_user_auth_passing_sample.agent import create_agent
+from google.adk.sessions import Session
+
+agent = create_agent()
+session = Session(
+    id="your_session_id",
+    app_name="your_app_name", 
+    user_id="your_user_id"
+)
+
+# Set user token in session state
+session.state['user_token'] = 'YOUR_ACTUAL_TOKEN_HERE'
+session.state['api_endpoint'] = 'https://your-internal-api.com'
+
+# Then use the agent in your workflow...
+```
+
+## Flow Diagram
+
+```mermaid
+graph TD
+    subgraph "User Application"
+        U[User]
+    end
+
+    subgraph "Agent Process"
+        A[Agent Instance<br/>per user-app-agentid]
+        S[Session State<br/>user_token, api_endpoint]
+        C[ContextToEnvMapperCallback]
+    end
+
+    subgraph "MCP Process"
+        M[MCP Server<br/>stdio transport]
+        E[Environment Variables<br/>USER_TOKEN, API_ENDPOINT]
+        API[Internal API Calls]
+    end
+
+    U -->|Sends request| A
+    A -->|Reads state| S
+    S -->|Extracts tokens| C
+    C -->|Maps to env vars| E
+    A -->|Spawns with env| M
+    M -->|Uses env vars| API
+    API -->|Response| M
+    M -->|Tool result| A
+    A -->|Response| U
+```
+
+## Context
+- Each agent instance is initiated per user-app-agentid.
+- The agent receives a user context (with token) and calls the MCP using stdio transport.
+- The MCP, built by the same organization, uses the token for internal API calls.
+- The ADK's context-to-env mapping feature makes this seamless.

contributing/samples/mcp_stdio_user_auth_passing_sample/agent.py

@@ -0,0 +1,60 @@
+"""
+Sample: Using ContextToEnvMapperCallback to pass user token from agent state to MCP via stdio transport.
+"""
+
+import os
+import tempfile
+from typing import Any
+from typing import Dict
+
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams
+from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset
+from mcp import StdioServerParameters
+
+_allowed_path = os.path.dirname(os.path.abspath(__file__))
+
+
+def user_token_env_mapper(state: Dict[str, Any]) -> Dict[str, str]:
+  """Extracts USER_TOKEN from agent state and maps to MCP env."""
+  env = {}
+  if "user_token" in state:
+    env["USER_TOKEN"] = state["user_token"]
+  if "api_endpoint" in state:
+    env["API_ENDPOINT"] = state["api_endpoint"]
+
+  print(f"Environment variables being passed to MCP: {env}")
+  return env
+
+
+def create_agent() -> LlmAgent:
+  """Create the agent with context to env mapper callback."""
+  # Create a temporary directory for the filesystem server
+  temp_dir = tempfile.mkdtemp()
+
+  return LlmAgent(
+      model="gemini-2.0-flash",
+      name="user_token_agent",
+      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}
+        """,
+      tools=[
+          MCPToolset(
+              connection_params=StdioConnectionParams(
+                  server_params=StdioServerParameters(
+                      command="npx",
+                      args=[
+                          "-y",  # Arguments for the command
+                          "@modelcontextprotocol/server-filesystem",
+                          _allowed_path,
+                      ],
+                  ),
+                  timeout=5,
+              ),
+              get_env_from_context_fn=user_token_env_mapper,
+              tool_filter=["read_file", "list_directory"],
+          )
+      ],
+  )

contributing/samples/mcp_stdio_user_auth_passing_sample/main.py

@@ -0,0 +1,95 @@
+"""
+Sample: Using ContextToEnvMapperCallback to pass user token from agent state to MCP via stdio transport.
+"""
+
+import asyncio
+
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.readonly_context import ReadonlyContext
+from google.adk.sessions import InMemorySessionService
+from google.adk.sessions import Session
+
+from .agent import create_agent
+
+
+async def main():
+  """Example of how to set up and run the agent with user token."""
+  print("=== STDIO MCP User Auth Passing Sample ===")
+  print()
+
+  # Create the agent
+  agent = create_agent()
+  print(f"✓ Created agent: {agent.name}")
+
+  # Create session service and session
+  session_service = InMemorySessionService()
+  session = Session(
+      id="sample_session",
+      app_name="stdio_mcp_user_auth_passing_sample",
+      user_id="sample_user",
+  )
+  print(f"✓ Created session: {session.id}")
+
+  # Set user token in session state
+  session.state["user_token"] = "sample_user_token_123"
+  session.state["api_endpoint"] = "https://internal-api.company.com"
+  print(f"✓ Set session state with user_token: {session.state['user_token']}")
+
+  # Create invocation context
+  invocation_context = InvocationContext(
+      invocation_id="sample_invocation",
+      agent=agent,
+      session=session,
+      session_service=session_service,
+  )
+
+  # Create readonly context
+  readonly_context = ReadonlyContext(invocation_context)
+  print(f"✓ Created readonly context")
+
+  print()
+  print("=== Demonstrating User Auth Token Passing to MCP ===")
+  print(
+      "Note: This sample shows how the callback extracts environment variables."
+  )
+  print("In a real scenario, these would be passed to an actual MCP server.")
+  print()
+
+  # Access the MCP toolset to demonstrate the callback
+  mcp_toolset = agent.tools[0]
+  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)
+
+    print(f"✓ Extracted environment variables:")
+    for key, value in env_vars.items():
+      print(f"   {key}={value}")
+    print()
+
+    print(
+        "✓ These environment variables would be injected into the MCP process"
+    )
+    print("✓ The MCP server can then use them for internal API calls")
+  else:
+    print("✗ No context-to-env mapper callback configured")
+
+  print()
+  print("=== Sample completed successfully! ===")
+  print()
+  print("Key points demonstrated:")
+  print("1. Session state holds user tokens and configuration")
+  print(
+      "2. Context-to-env mapper callback extracts these as environment"
+      " variables"
+  )
+  print("3. Environment variables would be passed to MCP server processes")
+  print("4. MCP servers can use these for authenticated API calls")
+
+
+if __name__ == "__main__":
+  asyncio.run(main())

pyproject.toml

@@ -38,7 +38,7 @@ dependencies = [
   "google-cloud-storage>=2.18.0, <3.0.0",           # For GCS Artifact service
   "google-genai>=1.21.1",                           # Google GenAI SDK
   "graphviz>=0.20.2",                               # Graphviz for graph rendering
-  "mcp>=1.8.0;python_version>='3.10'",              # For MCP Toolset
+  "mcp>=1.9.4;python_version>='3.10'",              # For MCP Toolset
   "opentelemetry-api>=1.31.0",                      # OpenTelemetry
   "opentelemetry-exporter-gcp-trace>=1.9.0",
   "opentelemetry-sdk>=1.31.0",

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

@@ -23,6 +23,7 @@
 import logging
 import sys
 from typing import Any
+from typing import Callable
 from typing import Dict
 from typing import Optional
 from typing import TextIO
@@ -177,13 +178,43 @@ def __init__(
     else:
       self._connection_params = connection_params
     self._errlog = errlog
-
     # Session pool: maps session keys to (session, exit_stack) tuples
     self._sessions: Dict[str, tuple[ClientSession, AsyncExitStack]] = {}
-
     # Lock to prevent race conditions in session creation
     self._session_lock = asyncio.Lock()
 
+  def update_connection_params(
+      self,
+      new_connection_params: Union[
+          StdioServerParameters,
+          StdioConnectionParams,
+          SseConnectionParams,
+          StreamableHTTPConnectionParams,
+      ],
+  ) -> None:
+    """Updates the connection parameters and invalidates existing sessions.
+
+    Args:
+        new_connection_params: New connection parameters to use.
+    """
+    if isinstance(new_connection_params, StdioServerParameters):
+      logger.warning(
+          'StdioServerParameters is not recommended. Please use'
+          ' StdioConnectionParams.'
+      )
+      self._connection_params = StdioConnectionParams(
+          server_params=new_connection_params,
+          timeout=5,
+      )
+    else:
+      self._connection_params = new_connection_params
+
+    # 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
+
+
   def _generate_session_key(
       self, merged_headers: Optional[Dict[str, str]] = None
   ) -> str: