From 86f05f37472cb5e85b1bf6f1b2787e02c215124c Mon Sep 17 00:00:00 2001 From: CorgiBoyG <111257566+CorgiBoyG@users.noreply.github.com> Date: Thu, 15 Oct 2026 15:52:20 +0800 Subject: [PATCH] fix(mcp): align auth and provider header precedence --- src/google/adk/tools/mcp_tool/mcp_tool.py | 29 ++++++-- src/google/adk/tools/mcp_tool/mcp_toolset.py | 17 ++--- .../unittests/tools/mcp_tool/test_mcp_tool.py | 72 +++++++++++++++++++ .../tools/mcp_tool/test_mcp_toolset.py | 32 +++++++++ 4 files changed, 134 insertions(+), 16 deletions(-) diff --git a/src/google/adk/tools/mcp_tool/mcp_tool.py b/src/google/adk/tools/mcp_tool/mcp_tool.py index 2b8b99fbeeb..ef147c4a91e 100644 --- a/src/google/adk/tools/mcp_tool/mcp_tool.py +++ b/src/google/adk/tools/mcp_tool/mcp_tool.py @@ -91,6 +91,24 @@ _UNSET = object() +def _merge_auth_and_provider_headers( + *, + auth_headers: dict[str, str] | None, + provider_headers: dict[str, str] | None, +) -> dict[str, str]: + """Merges MCP session headers, with explicit auth taking precedence.""" + headers = dict(provider_headers or {}) + if auth_headers: + auth_header_names = {name.lower() for name in auth_headers} + headers = { + name: value + for name, value in headers.items() + if name.lower() not in auth_header_names + } + headers.update(auth_headers) + return headers + + # Values the server owns: free-form JSON, not model fields. A `_meta` inside # one is the server's own -- a JSON Schema may even declare a property by that # name -- so the walk below must not descend into them. @@ -651,11 +669,10 @@ async def _run_async_impl( if inspect.isawaitable(dynamic_headers): dynamic_headers = await dynamic_headers - headers: dict[str, str] = {} - if auth_headers: - headers.update(auth_headers) - if dynamic_headers: - headers.update(dynamic_headers) + headers = _merge_auth_and_provider_headers( + auth_headers=auth_headers, + provider_headers=cast(dict[str, str] | None, dynamic_headers), + ) final_headers = headers if headers else None # Propagate trace context in the _meta field as sprcified by MCP protocol. @@ -842,7 +859,7 @@ async def _get_headers( """ headers: dict[str, str] | None = None if credential: - if credential.oauth2: + if credential.oauth2 and credential.oauth2.access_token: headers = {"Authorization": f"Bearer {credential.oauth2.access_token}"} elif credential.http: # Handle HTTP authentication schemes diff --git a/src/google/adk/tools/mcp_tool/mcp_toolset.py b/src/google/adk/tools/mcp_tool/mcp_toolset.py index abbc4b8a9a2..61068998e0a 100644 --- a/src/google/adk/tools/mcp_tool/mcp_toolset.py +++ b/src/google/adk/tools/mcp_tool/mcp_toolset.py @@ -24,6 +24,7 @@ from typing import Any from typing import Awaitable from typing import Callable +from typing import cast from typing import Dict from typing import List from typing import Optional @@ -63,6 +64,7 @@ from .mcp_session_manager import StreamableHTTPConnectionParams # MCPTool stays imported so existing mcp_toolset.MCPTool callers still resolve. from .mcp_tool import _dump_mcp_model +from .mcp_tool import _merge_auth_and_provider_headers from .mcp_tool import _RESERVED_TOOL_NAMES from .mcp_tool import MCPTool # pylint: disable=unused-import # noqa: F401 from .mcp_tool import McpTool @@ -357,22 +359,17 @@ async def _build_headers( Returns: The merged header provider and auth headers, empty if there are none. """ - headers: Dict[str, str] = {} - - # Add headers from header_provider if available + provider_headers = None 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) - # Add auth headers from exchanged credential if available auth_headers = self._get_auth_headers(readonly_context) - if auth_headers: - headers.update(auth_headers) - - return headers + return _merge_auth_and_provider_headers( + auth_headers=auth_headers, + provider_headers=cast(Dict[str, str] | None, provider_headers), + ) async def _execute_with_session( self, diff --git a/tests/unittests/tools/mcp_tool/test_mcp_tool.py b/tests/unittests/tools/mcp_tool/test_mcp_tool.py index 4c8274a40d7..7775146cba3 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_tool.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_tool.py @@ -1855,6 +1855,78 @@ async def test_run_async_impl_with_header_provider_and_oauth2(self): "test_tool", arguments=args, progress_callback=None, meta=None ) + @pytest.mark.asyncio + async def test_run_async_impl_auth_overrides_duplicate_provider_header(self): + """Explicit auth is authoritative when provider headers overlap.""" + header_provider = Mock( + return_value={ + "Authorization": "Bearer provider-token", + "X-Tenant-ID": "test-tenant", + } + ) + tool = MCPTool( + mcp_tool=self.mock_mcp_tool, + mcp_session_manager=self.mock_session_manager, + header_provider=header_provider, + ) + credential = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth(access_token="auth-token"), + ) + mcp_response = CallToolResult( + content=[TextContent(type="text", text="success")] + ) + self.mock_session.call_tool = AsyncMock(return_value=mcp_response) + tool_context = Mock(spec=ToolContext) + tool_context._invocation_context = Mock() + + await tool._run_async_impl( + args={"param1": "test_value"}, + tool_context=tool_context, + credential=credential, + ) + + self.mock_session_manager.create_session.assert_awaited_once_with( + headers={ + "Authorization": "Bearer auth-token", + "X-Tenant-ID": "test-tenant", + } + ) + + @pytest.mark.asyncio + async def test_run_async_impl_provider_survives_unusable_oauth_credential( + self, + ): + """An OAuth credential without a token must not replace provider auth.""" + header_provider = Mock( + return_value={"Authorization": "Bearer provider-token"} + ) + tool = MCPTool( + mcp_tool=self.mock_mcp_tool, + mcp_session_manager=self.mock_session_manager, + header_provider=header_provider, + ) + credential = AuthCredential( + auth_type=AuthCredentialTypes.OAUTH2, + oauth2=OAuth2Auth(access_token=None), + ) + mcp_response = CallToolResult( + content=[TextContent(type="text", text="success")] + ) + self.mock_session.call_tool = AsyncMock(return_value=mcp_response) + tool_context = Mock(spec=ToolContext) + tool_context._invocation_context = Mock() + + await tool._run_async_impl( + args={"param1": "test_value"}, + tool_context=tool_context, + credential=credential, + ) + + self.mock_session_manager.create_session.assert_awaited_once_with( + headers={"Authorization": "Bearer provider-token"} + ) + def test_init_with_progress_callback(self): """Test initialization with progress_callback.""" diff --git a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py index 6f8360d8062..5f7c4272e08 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_toolset.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_toolset.py @@ -530,6 +530,38 @@ async def header_provider(_context): headers=expected_headers ) + @pytest.mark.asyncio + async def test_get_tools_auth_overrides_duplicate_provider_header(self): + """Explicit auth is authoritative when provider headers overlap.""" + self.mock_session.list_tools = AsyncMock( + return_value=MockListToolsResult([MockMCPTool("tool1")]) + ) + mock_readonly_context = Mock(spec=ReadonlyContext) + toolset = McpToolset( + connection_params=self.mock_stdio_params, + header_provider=Mock( + return_value={ + "authorization": "Bearer provider-token", + "X-Tenant-ID": "test-tenant", + } + ), + ) + toolset._mcp_session_manager = self.mock_session_manager + + with patch.object( + toolset, + "_get_auth_headers", + return_value={"Authorization": "Bearer auth-token"}, + ): + await toolset.get_tools(readonly_context=mock_readonly_context) + + self.mock_session_manager.create_session.assert_awaited_once_with( + headers={ + "Authorization": "Bearer auth-token", + "X-Tenant-ID": "test-tenant", + } + ) + @pytest.mark.asyncio async def test_close_success(self): """Test successful cleanup."""