feat: Support dynamic per-request headers in MCPToolset

Add a header_provider param which is a callable[ReadonlyContext, Dict[str, Any]] for users to build headers in MCPToolset
fix: #3156
PiperOrigin-RevId: 820412372

Author: wukath

contributing/samples/mcp_dynamic_header_agent/README.md

@@ -0,0 +1,8 @@
+This agent connects to a local MCP server via Streamable HTTP and provides
+custom per-request headers to the MCP server.
+
+To run this agent, start the local MCP server first by running:
+
+```bash
+uv run header_server.py
+```

contributing/samples/mcp_dynamic_header_agent/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from . import agent

contributing/samples/mcp_dynamic_header_agent/agent.py

@@ -0,0 +1,34 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
+from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
+
+root_agent = LlmAgent(
+    model='gemini-2.0-flash',
+    name='tenant_agent',
+    instruction="""You are a helpful assistant that helps users get tenant
+ information. Call the get_tenant_data tool when the user asks for tenant data.""",
+    tools=[
+        McpToolset(
+            connection_params=StreamableHTTPConnectionParams(
+                url='http://localhost:3000/mcp',
+            ),
+            tool_filter=['get_tenant_data'],
+            header_provider=lambda ctx: {'X-Tenant-ID': 'tenant1'},
+        )
+    ],
+)

contributing/samples/mcp_dynamic_header_agent/header_server.py

@@ -0,0 +1,50 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from fastapi import Request
+from mcp.server.fastmcp import Context
+from mcp.server.fastmcp import FastMCP
+
+mcp = FastMCP('Header Check Server', host='localhost', port=3000)
+
+TENANT_DATA = {
+    'tenant1': {'name': 'Tenant 1', 'data': 'Data for tenant 1'},
+    'tenant2': {'name': 'Tenant 2', 'data': 'Data for tenant 2'},
+}
+
+
+@mcp.tool(
+    description='Returns tenant specific data based on X-Tenant-ID header.'
+)
+def get_tenant_data(context: Context) -> dict:
+  """Return tenant specific data."""
+  if context.request_context and context.request_context.request:
+    headers = context.request_context.request.headers
+    tenant_id = headers.get('x-tenant-id')
+    if tenant_id in TENANT_DATA:
+      return TENANT_DATA[tenant_id]
+    else:
+      return {'error': f'Tenant {tenant_id} not found'}
+  else:
+    return {'error': 'Could not get request context'}
+
+
+if __name__ == '__main__':
+  try:
+    print('Starting Header Check MCP server on http://localhost:3000')
+    mcp.run(transport='streamable-http')
+  except KeyboardInterrupt:
+    print('\nServer stopped.')

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

@@ -17,8 +17,10 @@
 import base64
 import inspect
 import logging
+import sys
 from typing import Any
 from typing import Callable
+from typing import Dict
 from typing import Optional
 from typing import Union
 import warnings
@@ -27,6 +29,7 @@
 from google.genai.types import FunctionDeclaration
 from typing_extensions import override
 
+from ...agents.readonly_context import ReadonlyContext
 from .._gemini_schema_util import _to_gemini_schema
 from .mcp_session_manager import MCPSessionManager
 from .mcp_session_manager import retry_on_closed_resource
@@ -36,8 +39,6 @@
 try:
   from mcp.types import Tool as McpBaseTool
 except ImportError as e:
-  import sys
-
   if sys.version_info < (3, 10):
     raise ImportError(
         "MCP Tool requires Python 3.10 or above. Please upgrade your Python"
@@ -75,6 +76,9 @@ def __init__(
       auth_scheme: Optional[AuthScheme] = None,
       auth_credential: Optional[AuthCredential] = None,
       require_confirmation: Union[bool, Callable[..., bool]] = False,
+      header_provider: Optional[
+          Callable[[ReadonlyContext], Dict[str, str]]
+      ] = None,
   ):
     """Initializes an MCPTool.
 
@@ -106,6 +110,7 @@ def __init__(
     self._mcp_tool = mcp_tool
     self._mcp_session_manager = mcp_session_manager
     self._require_confirmation = require_confirmation
+    self._header_provider = header_provider
 
   @override
   def _get_declaration(self) -> FunctionDeclaration:
@@ -192,10 +197,24 @@ async def _run_async_impl(
         Any: The response from the tool.
     """
     # Extract headers from credential for session pooling
-    headers = await self._get_headers(tool_context, credential)
+    auth_headers = await self._get_headers(tool_context, credential)
+    dynamic_headers = None
+    if self._header_provider:
+      dynamic_headers = self._header_provider(
+          ReadonlyContext(tool_context._invocation_context)
+      )
+
+    headers: Dict[str, str] = {}
+    if auth_headers:
+      headers.update(auth_headers)
+    if dynamic_headers:
+      headers.update(dynamic_headers)
+    final_headers = headers if headers else None
 
     # Get the session from the session manager
-    session = await self._mcp_session_manager.create_session(headers=headers)
+    session = await self._mcp_session_manager.create_session(
+        headers=final_headers
+    )
 
     response = await session.call_tool(self._mcp_tool.name, arguments=args)
     return response

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

@@ -16,6 +16,8 @@
 
 import logging
 import sys
+from typing import Any
+from typing import AsyncIterator
 from typing import Callable
 from typing import Dict
 from typing import List
@@ -107,6 +109,9 @@ def __init__(
       auth_scheme: Optional[AuthScheme] = None,
       auth_credential: Optional[AuthCredential] = None,
       require_confirmation: Union[bool, Callable[..., bool]] = False,
+      header_provider: Optional[
+          Callable[[ReadonlyContext], Dict[str, str]]
+      ] = None,
   ):
     """Initializes the MCPToolset.
 
@@ -130,6 +135,8 @@ def __init__(
       require_confirmation: Whether tools in this toolset require
         confirmation. Can be a single boolean or a callable to apply to all
         tools.
+      header_provider: A callable that takes a ReadonlyContext and returns a
+        dictionary of headers to be used for the MCP session.
     """
     super().__init__(tool_filter=tool_filter, tool_name_prefix=tool_name_prefix)
 
@@ -138,6 +145,7 @@ def __init__(
 
     self._connection_params = connection_params
     self._errlog = errlog
+    self._header_provider = header_provider
 
     # Create the session manager that will handle the MCP connection
     self._mcp_session_manager = MCPSessionManager(
@@ -162,8 +170,13 @@ async def get_tools(
     Returns:
         List[BaseTool]: A list of tools available under the specified context.
     """
+    headers = (
+        self._header_provider(readonly_context)
+        if self._header_provider and readonly_context
+        else None
+    )
     # Get session from session manager
-    session = await self._mcp_session_manager.create_session()
+    session = await self._mcp_session_manager.create_session(headers=headers)
 
     # Fetch available tools from the MCP server
     tools_response: ListToolsResult = await session.list_tools()
@@ -177,6 +190,7 @@ async def get_tools(
           auth_scheme=self._auth_scheme,
           auth_credential=self._auth_credential,
           require_confirmation=self._require_confirmation,
+          header_provider=self._header_provider,
       )
 
       if self._is_tool_selected(mcp_tool, readonly_context):

tests/unittests/tools/mcp_tool/test_mcp_tool.py

@@ -640,3 +640,74 @@ def test_init_validation(self):
 
     with pytest.raises(TypeError):
       MCPTool(mcp_tool=self.mock_mcp_tool)  # Missing session manager
+
+  @pytest.mark.asyncio
+  async def test_run_async_impl_with_header_provider_no_auth(self):
+    """Test running tool with header_provider but no auth."""
+    expected_headers = {"X-Tenant-ID": "test-tenant"}
+    header_provider = Mock(return_value=expected_headers)
+    tool = MCPTool(
+        mcp_tool=self.mock_mcp_tool,
+        mcp_session_manager=self.mock_session_manager,
+        header_provider=header_provider,
+    )
+
+    expected_response = {"result": "success"}
+    self.mock_session.call_tool = AsyncMock(return_value=expected_response)
+
+    tool_context = Mock(spec=ToolContext)
+    tool_context._invocation_context = Mock()
+    args = {"param1": "test_value"}
+
+    result = await tool._run_async_impl(
+        args=args, tool_context=tool_context, credential=None
+    )
+
+    assert result == expected_response
+    header_provider.assert_called_once()
+    self.mock_session_manager.create_session.assert_called_once_with(
+        headers=expected_headers
+    )
+    self.mock_session.call_tool.assert_called_once_with(
+        "test_tool", arguments=args
+    )
+
+  @pytest.mark.asyncio
+  async def test_run_async_impl_with_header_provider_and_oauth2(self):
+    """Test running tool with header_provider and OAuth2 auth."""
+    dynamic_headers = {"X-Tenant-ID": "test-tenant"}
+    header_provider = Mock(return_value=dynamic_headers)
+    tool = MCPTool(
+        mcp_tool=self.mock_mcp_tool,
+        mcp_session_manager=self.mock_session_manager,
+        header_provider=header_provider,
+    )
+
+    oauth2_auth = OAuth2Auth(access_token="test_access_token")
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2, oauth2=oauth2_auth
+    )
+
+    expected_response = {"result": "success"}
+    self.mock_session.call_tool = AsyncMock(return_value=expected_response)
+
+    tool_context = Mock(spec=ToolContext)
+    tool_context._invocation_context = Mock()
+    args = {"param1": "test_value"}
+
+    result = await tool._run_async_impl(
+        args=args, tool_context=tool_context, credential=credential
+    )
+
+    assert result == expected_response
+    header_provider.assert_called_once()
+    self.mock_session_manager.create_session.assert_called_once()
+    call_args = self.mock_session_manager.create_session.call_args
+    headers = call_args[1]["headers"]
+    assert headers == {
+        "Authorization": "Bearer test_access_token",
+        "X-Tenant-ID": "test-tenant",
+    }
+    self.mock_session.call_tool.assert_called_once_with(
+        "test_tool", arguments=args
+    )

tests/unittests/tools/mcp_tool/test_mcp_toolset.py

@@ -29,6 +29,7 @@
 
 # Import dependencies with version checking
 try:
+  from google.adk.agents.readonly_context import ReadonlyContext
   from google.adk.tools.mcp_tool.mcp_session_manager import MCPSessionManager
   from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams
   from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams
@@ -55,6 +56,7 @@ def __init__(self, command="test_command", args=None):
     StreamableHTTPConnectionParams = DummyClass
     MCPTool = DummyClass
     MCPToolset = DummyClass
+    ReadonlyContext = DummyClass
   else:
     raise e
 
@@ -245,6 +247,31 @@ def file_tools_filter(tool, context):
     assert tools[0].name == "read_file"
     assert tools[1].name == "write_file"
 
+  @pytest.mark.asyncio
+  async def test_get_tools_with_header_provider(self):
+    """Test get_tools with a header_provider."""
+    mock_tools = [MockMCPTool("tool1"), MockMCPTool("tool2")]
+    self.mock_session.list_tools = AsyncMock(
+        return_value=MockListToolsResult(mock_tools)
+    )
+    mock_readonly_context = Mock(spec=ReadonlyContext)
+    expected_headers = {"X-Tenant-ID": "test-tenant"}
+    header_provider = Mock(return_value=expected_headers)
+
+    toolset = MCPToolset(
+        connection_params=self.mock_stdio_params,
+        header_provider=header_provider,
+    )
+    toolset._mcp_session_manager = self.mock_session_manager
+
+    tools = await toolset.get_tools(readonly_context=mock_readonly_context)
+
+    assert len(tools) == 2
+    header_provider.assert_called_once_with(mock_readonly_context)
+    self.mock_session_manager.create_session.assert_called_once_with(
+        headers=expected_headers
+    )
+
   @pytest.mark.asyncio
   async def test_close_success(self):
     """Test successful cleanup."""