fix(mcp): await async MCP header providers

Merge #6105

## Summary
- await async/awaitable MCP header providers in async execution paths.
- add unit tests verifying the async header provider functionality.

Fixes #6090

Co-authored-by: Kathy Wu <wukathy@google.com>
COPYBARA_INTEGRATE_REVIEW=#6105 from he-yufeng:fix/await-async-header-provider 086d352
PiperOrigin-RevId: 941317964

Author: 2 people

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

@@ -16,16 +16,13 @@
 
 import asyncio
 import base64
+from collections.abc import Awaitable
 import inspect
 import logging
 from typing import Any
 from typing import Callable
-from typing import Dict
-from typing import List
-from typing import Optional
 from typing import Protocol
 from typing import runtime_checkable
-from typing import Union
 import warnings
 
 from fastapi.openapi.models import APIKeyIn
@@ -104,9 +101,9 @@ def __call__(
       self,
       tool_name: str,
       *,
-      callback_context: Optional[CallbackContext] = None,
+      callback_context: CallbackContext | None = None,
       **kwargs: Any,
-  ) -> Optional[ProgressFnT]:
+  ) -> ProgressFnT | None:
     """Create a progress callback for a specific tool.
 
     Args:
@@ -139,15 +136,17 @@ def __init__(
       *,
       mcp_tool: McpBaseTool,
       mcp_session_manager: MCPSessionManager,
-      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,
-      progress_callback: Optional[
-          Union[ProgressFnT, ProgressCallbackFactory]
-      ] = None,
+      auth_scheme: AuthScheme | None = None,
+      auth_credential: AuthCredential | None = None,
+      require_confirmation: bool | Callable[..., bool] = False,
+      header_provider: (
+          Callable[
+              [ReadonlyContext],
+              dict[str, str] | Awaitable[dict[str, str]],
+          ]
+          | None
+      ) = None,
+      progress_callback: ProgressFnT | ProgressCallbackFactory | None = None,
   ):
     """Initializes an McpTool.
 
@@ -225,7 +224,7 @@ def raw_mcp_tool(self) -> McpBaseTool:
     return self._mcp_tool
 
   @property
-  def visibility(self) -> List[str]:
+  def visibility(self) -> list[str]:
     """Returns the visibility if this MCP tool meta has one."""
     meta = getattr(self.raw_mcp_tool, "meta", None)
     if not meta or not isinstance(meta, dict):
@@ -238,7 +237,7 @@ def visibility(self) -> List[str]:
     return []
 
   @property
-  def mcp_app_resource_uri(self) -> Optional[str]:
+  def mcp_app_resource_uri(self) -> str | None:
     """Returns the MCP App UI resource URI if this tool has one.
 
     MCP Apps declare a UI resource via `meta.ui.resourceUri` in the tool
@@ -379,7 +378,7 @@ async def run_async(
   @override
   async def _run_async_impl(
       self, *, args, tool_context: ToolContext, credential: AuthCredential
-  ) -> Dict[str, Any]:
+  ) -> dict[str, Any]:
     """Runs the tool asynchronously.
 
     Args:
@@ -396,8 +395,10 @@ async def _run_async_impl(
       dynamic_headers = self._header_provider(
           ReadonlyContext(tool_context._invocation_context)  # pylint: disable=protected-access
       )
+      if inspect.isawaitable(dynamic_headers):
+        dynamic_headers = await dynamic_headers
 
-    headers: Dict[str, str] = {}
+    headers: dict[str, str] = {}
     if auth_headers:
       headers.update(auth_headers)
     if dynamic_headers:
@@ -406,7 +407,7 @@ async def _run_async_impl(
 
     # Propagate trace context in the _meta field as sprcified by MCP protocol.
     # See https://agentclientprotocol.com/protocol/extensibility#the-meta-field
-    trace_carrier: Dict[str, str] = {}
+    trace_carrier: dict[str, str] = {}
     propagate.get_global_textmap().inject(carrier=trace_carrier)
     meta_trace_context = trace_carrier if trace_carrier else None
 
@@ -468,15 +469,15 @@ async def _run_async_impl(
       )
     return result
 
-  def _detect_error_in_response(self, response: Any) -> Optional[str]:
+  def _detect_error_in_response(self, response: Any) -> str | None:
     """Telemetry hook: returns an error type if the response indicates an error."""
     if isinstance(response, dict) and response.get("isError"):
       return "MCP_TOOL_ERROR"
     return None
 
   def _resolve_progress_callback(
       self, tool_context: ToolContext
-  ) -> Optional[ProgressFnT]:
+  ) -> ProgressFnT | None:
     """Resolve the progress callback for the current invocation.
 
     If progress_callback is a ProgressCallbackFactory, call it to create
@@ -510,7 +511,7 @@ def _resolve_progress_callback(
 
   async def _get_headers(
       self, tool_context: ToolContext, credential: AuthCredential
-  ) -> Optional[dict[str, str]]:
+  ) -> dict[str, str] | None:
     """Extracts authentication headers from credentials.
 
     Args:
@@ -524,7 +525,7 @@ async def _get_headers(
         ValueError: If API key authentication is configured for non-header
         location.
     """
-    headers: Optional[dict[str, str]] = None
+    headers: dict[str, str] | None = None
     if credential:
       if credential.oauth2:
         headers = {"Authorization": f"Bearer {credential.oauth2.access_token}"}

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

@@ -16,6 +16,7 @@
 
 import asyncio
 import base64
+import inspect
 import logging
 import sys
 from typing import Any
@@ -107,9 +108,13 @@ 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,
+      header_provider: (
+          Callable[
+              [ReadonlyContext],
+              dict[str, str] | Awaitable[dict[str, str]],
+          ]
+          | None
+      ) = None,
       progress_callback: Optional[
           Union[ProgressFnT, ProgressCallbackFactory]
       ] = None,
@@ -293,6 +298,8 @@ async def _execute_with_session(
     # Add headers from header_provider if available
     if self._header_provider and readonly_context:
       provider_headers = self._header_provider(readonly_context)
+      if inspect.isawaitable(provider_headers):
+        provider_headers = await provider_headers
       if provider_headers:
         headers.update(provider_headers)
 

tests/unittests/tools/mcp_tool/test_mcp_tool.py

@@ -953,6 +953,41 @@ async def test_run_async_impl_with_header_provider_no_auth(self):
         "test_tool", arguments=args, progress_callback=None, meta=None
     )
 
+  @pytest.mark.asyncio
+  async def test_run_async_impl_with_async_header_provider_no_auth(self):
+    """Test running tool with an async header_provider and no authentication."""
+    expected_headers = {"X-Tenant-ID": "test-tenant"}
+
+    async def header_provider(_context):
+      return expected_headers
+
+    tool = MCPTool(
+        mcp_tool=self.mock_mcp_tool,
+        mcp_session_manager=self.mock_session_manager,
+        header_provider=header_provider,
+    )
+
+    mcp_response = CallToolResult(
+        content=[TextContent(type="text", text="response text")]
+    )
+    self.mock_session.call_tool = AsyncMock(return_value=mcp_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 == mcp_response.model_dump(exclude_none=True, mode="json")
+    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, progress_callback=None, meta=None
+    )
+
   @pytest.mark.asyncio
   async def test_run_async_impl_with_header_provider_and_oauth2(self):
     """Test running tool with header_provider and OAuth2 auth."""

tests/unittests/tools/mcp_tool/test_mcp_toolset.py

@@ -303,6 +303,32 @@ async def test_get_tools_with_header_provider(self):
         headers=expected_headers
     )
 
+  @pytest.mark.asyncio
+  async def test_get_tools_with_async_header_provider(self):
+    """Test get_tools with an async 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"}
+
+    async def header_provider(_context):
+      return 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
+    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."""