From 055f539db07e42070deec4c787768d439267d56f Mon Sep 17 00:00:00 2001 From: Chaitanya Laxman Date: Mon, 7 Sep 2026 20:18:39 +0400 Subject: [PATCH] feat(mcp): propagate grounding metadata from MCP _meta _maybe_add_grounding_metadata returned early unless a tool was named google_search_agent. git log -S shows d689a04f (Xuan Yang) replaced that with getattr(tool, 'propagate_grounding_metadata', False), and b2daf83d (Haiyuan Cao) put the name check back. I ran _handle_after_model_callback with dummy_tool and temp:_adk_grounding_metadata={'foo': 'bar'}: result was None. Restore the d689a04f gate and the canonical_tools_cache read. Add propagate_grounding_metadata on McpTool and McpToolset so _meta.adk_grounding_metadata can populate the same temp key. Verification: 15 passed. upstream/main src: cache test AttributeError on None grounding_metadata; MCP flag-on test TypeError on the new kwarg. Fixes #6081 --- .../adk/flows/llm_flows/base_llm_flow.py | 9 ++- src/google/adk/tools/mcp_tool/mcp_tool.py | 34 +++++++++ src/google/adk/tools/mcp_tool/mcp_toolset.py | 6 ++ .../flows/llm_flows/test_base_llm_flow.py | 18 ++--- .../unittests/tools/mcp_tool/test_mcp_tool.py | 69 +++++++++++++++++++ 5 files changed, 125 insertions(+), 11 deletions(-) diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index 975f0e528ca..867f089ddba 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -275,8 +275,9 @@ async def _handle_after_model_callback( ) -> Optional[LlmResponse]: """Runs after-model callbacks (plugins then agent callbacks). - Also handles grounding metadata injection when google_search_agent is - among the agent's tools. + Also handles grounding metadata injection when a tool sets + ``propagate_grounding_metadata`` and ``temp:_adk_grounding_metadata`` + is present on the session. Args: invocation_context: The invocation context. @@ -298,7 +299,9 @@ async def _maybe_add_grounding_metadata( tools = await agent.canonical_tools(readonly_context) invocation_context.canonical_tools_cache = tools - if not any(tool.name == 'google_search_agent' for tool in tools): + if not any( + getattr(tool, 'propagate_grounding_metadata', False) for tool in tools + ): return response ground_metadata = invocation_context.session.state.get( 'temp:_adk_grounding_metadata', None diff --git a/src/google/adk/tools/mcp_tool/mcp_tool.py b/src/google/adk/tools/mcp_tool/mcp_tool.py index 83c426875de..1a683d40b64 100644 --- a/src/google/adk/tools/mcp_tool/mcp_tool.py +++ b/src/google/adk/tools/mcp_tool/mcp_tool.py @@ -27,7 +27,9 @@ from fastapi.openapi.models import APIKeyIn from google.genai.types import FunctionDeclaration +from google.genai.types import GroundingMetadata from opentelemetry import propagate +from pydantic import ValidationError from typing_extensions import override from ...agents.callback_context import CallbackContext @@ -291,6 +293,7 @@ def __init__( | None ) = None, progress_callback: ProgressFnT | ProgressCallbackFactory | None = None, + propagate_grounding_metadata: bool = False, ): """Initializes an McpTool. @@ -317,6 +320,10 @@ def __init__( The factory receives (tool_name, callback_context, **kwargs) and returns a ProgressFnT or None. This allows callbacks to access and modify runtime context like session state. + propagate_grounding_metadata: If True, copy + ``meta.adk_grounding_metadata`` from the MCP result into + ``temp:_adk_grounding_metadata`` so the flow can attach it to + ``LlmResponse``. Default False. Raises: ValueError: If the MCP tool name collides with a reserved ADK tool @@ -342,6 +349,7 @@ def __init__( self._require_confirmation = require_confirmation self._header_provider = header_provider self._progress_callback = progress_callback + self.propagate_grounding_metadata = propagate_grounding_metadata @override def _get_declaration(self) -> FunctionDeclaration: @@ -634,6 +642,7 @@ async def _run_async_impl( # Keep the caller's key names off the installed SDK's field naming. result = _dump_mcp_model(response) + self._store_grounding_metadata_from_result(result, tool_context) # 2.x-only field. Acting on it (`input_required` drives elicitation) is a # feature, not compatibility. Not dropped on 1.x, where a key of that name @@ -664,6 +673,31 @@ async def _run_async_impl( ) return result + def _store_grounding_metadata_from_result( + self, result: dict[str, Any], tool_context: ToolContext + ) -> None: + """Copies ADK grounding from MCP meta into session temp state.""" + if not self.propagate_grounding_metadata: + return + meta = result.get("meta") + if meta is None: + meta = result.get("_meta") + if not isinstance(meta, dict): + return + raw = meta.get("adk_grounding_metadata") + if raw is None: + return + try: + metadata = GroundingMetadata.model_validate(raw) + except ValidationError as e: + logger.warning( + "Ignoring _meta.adk_grounding_metadata from %s: %s", + self.name, + e, + ) + return + tool_context.state["temp:_adk_grounding_metadata"] = metadata + def _detect_error_in_response(self, response: Any) -> str | None: """Telemetry hook: returns an error type if the response indicates an error.""" # `response` is a dumped CallToolResult. `_run_async_impl` restores diff --git a/src/google/adk/tools/mcp_tool/mcp_toolset.py b/src/google/adk/tools/mcp_tool/mcp_toolset.py index dcf4a7e0c68..392b4626447 100644 --- a/src/google/adk/tools/mcp_tool/mcp_toolset.py +++ b/src/google/adk/tools/mcp_tool/mcp_toolset.py @@ -171,6 +171,7 @@ def __init__( sampling_capabilities: SamplingCapability | None = None, elicitation_callback: ElicitationFnT | None = None, credential_key: str | None = None, + propagate_grounding_metadata: bool = False, ): """Initializes the McpToolset. @@ -223,6 +224,9 @@ def __init__( elicitations used for out-of-band flows such as auth challenges. credential_key: A user specified key used to load and save this credential in a credential service. Used with auth_scheme. + propagate_grounding_metadata: If True, each listed tool copies + ``meta.adk_grounding_metadata`` from the MCP result into + ``temp:_adk_grounding_metadata``. Default False. """ super().__init__(tool_filter=tool_filter, tool_name_prefix=tool_name_prefix) @@ -264,6 +268,7 @@ def __init__( self._auth_scheme = auth_scheme self._auth_credential = auth_credential self._require_confirmation = require_confirmation + self._propagate_grounding_metadata = propagate_grounding_metadata # Store auth config as instance variable so ADK can populate # exchanged_auth_credential in-place before calling get_tools() self._auth_config: Optional[AuthConfig] = ( @@ -532,6 +537,7 @@ async def get_tools( progress_callback=self._progress_callback if hasattr(self, "_progress_callback") else None, + propagate_grounding_metadata=self._propagate_grounding_metadata, ) if self._is_tool_selected(mcp_tool, readonly_context): diff --git a/tests/unittests/flows/llm_flows/test_base_llm_flow.py b/tests/unittests/flows/llm_flows/test_base_llm_flow.py index 37149bd11e6..7ab1ded4256 100644 --- a/tests/unittests/flows/llm_flows/test_base_llm_flow.py +++ b/tests/unittests/flows/llm_flows/test_base_llm_flow.py @@ -611,6 +611,9 @@ async def test_handle_after_model_callback_grounding_with_callback_override( agent_response.grounding_metadata = state_metadata assert result == agent_response + assert result.grounding_metadata == ( + state_metadata if expect_metadata else None + ) agent_callback.assert_called_once() @@ -672,6 +675,9 @@ def __init__(self): plugin_response.grounding_metadata = state_metadata assert result == plugin_response + assert result.grounding_metadata == ( + state_metadata if expect_metadata else None + ) plugin.after_model_callback.assert_called_once() @@ -685,16 +691,16 @@ async def mock_canonical_tools(self, readonly_context=None): canonical_tools_call_count += 1 from google.adk.tools.base_tool import BaseTool - class MockGoogleSearchTool(BaseTool): + class MockResearchTool(BaseTool): def __init__(self): - super().__init__(name='google_search_agent', description='Mock search') + super().__init__(name='research_agent', description='Mock research') self.propagate_grounding_metadata = True async def call(self, **kwargs): return 'mock result' - return [MockGoogleSearchTool()] + return [MockResearchTool()] agent = Agent(name='test_agent', tools=[google_search, dummy_tool]) @@ -720,7 +726,6 @@ async def call(self, **kwargs): author=agent.name, ) - # Call _handle_after_model_callback multiple times with the same context result1 = await _handle_after_model_callback( invocation_context, llm_response, event ) @@ -738,10 +743,7 @@ async def call(self, **kwargs): assert invocation_context.canonical_tools_cache is not None assert len(invocation_context.canonical_tools_cache) == 1 - assert ( - invocation_context.canonical_tools_cache[0].name - == 'google_search_agent' - ) + assert invocation_context.canonical_tools_cache[0].name == 'research_agent' assert result1.grounding_metadata == {'foo': 'bar'} assert result2.grounding_metadata == {'foo': 'bar'} diff --git a/tests/unittests/tools/mcp_tool/test_mcp_tool.py b/tests/unittests/tools/mcp_tool/test_mcp_tool.py index 245b22d2d61..01acdc889e9 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_tool.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_tool.py @@ -22,6 +22,8 @@ from unittest.mock import patch from google.adk.agents.context import Context +from google.adk.agents.invocation_context import InvocationContext +from google.adk.agents.llm_agent import Agent from google.adk.auth.auth_credential import AuthCredential from google.adk.auth.auth_credential import AuthCredentialTypes from google.adk.auth.auth_credential import HttpAuth @@ -33,6 +35,7 @@ from google.adk.events.event_actions import EventActions from google.adk.features import FeatureName from google.adk.features._feature_registry import temporary_feature_override +from google.adk.sessions.in_memory_session_service import InMemorySessionService from google.adk.tools.mcp_tool import mcp_tool from google.adk.tools.mcp_tool.mcp_session_manager import _SESSION_IDLE_TTL_SECONDS from google.adk.tools.mcp_tool.mcp_session_manager import MCPSessionManager @@ -42,6 +45,7 @@ from google.adk.tools.mcp_tool.mcp_tool import ProgressFnT from google.adk.tools.tool_context import ToolContext from google.genai.types import FunctionDeclaration +from google.genai.types import GroundingMetadata from mcp.types import CallToolResult from mcp.types import ImageContent from mcp.types import TextContent @@ -403,6 +407,71 @@ async def test_run_async_impl_no_auth(self): "test_tool", arguments=args, progress_callback=None, meta=None ) + async def _tool_context_with_session(self) -> ToolContext: + session_service = InMemorySessionService() + session = await session_service.create_session( + app_name="test_app", user_id="test_user" + ) + tool_context = ToolContext( + invocation_context=InvocationContext( + invocation_id="invocation_id", + agent=Agent(name="test_agent"), + session=session, + session_service=session_service, + ) + ) + tool_context.function_call_id = "test-call-id" + return tool_context + + @pytest.mark.asyncio + async def test_run_async_impl_propagates_grounding_metadata_from_meta(self): + """_meta.adk_grounding_metadata becomes temp state when the flag is on.""" + tool = MCPTool( + mcp_tool=self.mock_mcp_tool, + mcp_session_manager=self.mock_session_manager, + propagate_grounding_metadata=True, + ) + mcp_response = CallToolResult( + content=[TextContent(type="text", text="success")], + _meta={"adk_grounding_metadata": {"webSearchQueries": ["q1"]}}, + ) + self.mock_session.call_tool = AsyncMock(return_value=mcp_response) + tool_context = await self._tool_context_with_session() + + result = await tool._run_async_impl( + args={"param1": "test_value"}, + tool_context=tool_context, + credential=None, + ) + + assert result == expected_tool_result(mcp_response) + stored = tool_context.state["temp:_adk_grounding_metadata"] + assert isinstance(stored, GroundingMetadata) + assert stored.web_search_queries == ["q1"] + + @pytest.mark.asyncio + async def test_run_async_impl_skips_grounding_metadata_when_flag_off(self): + """Default McpTool leaves temp grounding unset even if _meta carries it.""" + tool = MCPTool( + mcp_tool=self.mock_mcp_tool, + mcp_session_manager=self.mock_session_manager, + ) + mcp_response = CallToolResult( + content=[TextContent(type="text", text="success")], + _meta={"adk_grounding_metadata": {"webSearchQueries": ["q1"]}}, + ) + self.mock_session.call_tool = AsyncMock(return_value=mcp_response) + tool_context = await self._tool_context_with_session() + + result = await tool._run_async_impl( + args={"param1": "test_value"}, + tool_context=tool_context, + credential=None, + ) + + assert result == expected_tool_result(mcp_response) + assert "temp:_adk_grounding_metadata" not in tool_context.state + @pytest.mark.asyncio async def test_in_flight_tool_call_is_held_out_of_the_idle_sweep(self): """A call in flight must not have its session swept out from under it."""