diff --git a/src/google/adk/tools/mcp_tool/mcp_tool.py b/src/google/adk/tools/mcp_tool/mcp_tool.py index 26e12bc1..4c06451c 100644 --- a/src/google/adk/tools/mcp_tool/mcp_tool.py +++ b/src/google/adk/tools/mcp_tool/mcp_tool.py @@ -16,16 +16,13 @@ from __future__ import annotations 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 @@ class ProgressCallbackFactory(Protocol): 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 @@ class McpTool(BaseAuthenticatedTool): *, 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 @@ class McpTool(BaseAuthenticatedTool): 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 @@ class McpTool(BaseAuthenticatedTool): 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 @@ class McpTool(BaseAuthenticatedTool): @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 @@ class McpTool(BaseAuthenticatedTool): 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 @@ class McpTool(BaseAuthenticatedTool): # 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,7 +469,7 @@ class McpTool(BaseAuthenticatedTool): ) 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" @@ -476,7 +477,7 @@ class McpTool(BaseAuthenticatedTool): 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 @@ class McpTool(BaseAuthenticatedTool): 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 @@ class McpTool(BaseAuthenticatedTool): 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}"} diff --git a/src/google/adk/tools/mcp_tool/mcp_toolset.py b/src/google/adk/tools/mcp_tool/mcp_toolset.py index faf615c7..d736b00c 100644 --- a/src/google/adk/tools/mcp_tool/mcp_toolset.py +++ b/src/google/adk/tools/mcp_tool/mcp_toolset.py @@ -16,6 +16,7 @@ from __future__ import annotations import asyncio import base64 +import inspect import logging import sys from typing import Any @@ -107,9 +108,13 @@ class McpToolset(BaseToolset): 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 @@ class McpToolset(BaseToolset): # 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) diff --git a/tests/unittests/tools/mcp_tool/test_mcp_tool.py b/tests/unittests/tools/mcp_tool/test_mcp_tool.py index 8d1af680..fefd04f1 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_tool.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_tool.py @@ -953,6 +953,41 @@ class TestMCPTool: "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.""" diff --git a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py index 20ec612e..0c1700ae 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py @@ -303,6 +303,32 @@ class TestMcpToolset: 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."""