feat: Support propagating grounding metadata from AgentTool

Close: #4671

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 891827874

Author: xuanyang15

contributing/samples/agent_tool_with_grounding_metadata/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2026 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/agent_tool_with_grounding_metadata/agent.py

@@ -0,0 +1,74 @@
+# Copyright 2026 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.
+
+import os
+import random
+
+from dotenv import load_dotenv
+from google.adk import Agent
+from google.adk.tools.agent_tool import AgentTool
+from google.adk.tools.tool_context import ToolContext
+from google.adk.tools.vertex_ai_search_tool import VertexAiSearchTool
+
+load_dotenv(override=True)
+
+VERTEXAI_DATASTORE_ID = os.getenv("VERTEXAI_DATASTORE_ID")
+if not VERTEXAI_DATASTORE_ID:
+  raise ValueError("VERTEXAI_DATASTORE_ID environment variable not set")
+
+
+def roll_die(sides: int, tool_context: ToolContext) -> int:
+  """Roll a die and return the rolled result.
+
+  Args:
+    sides: The integer number of sides the die has.
+
+  Returns:
+    An integer of the result of rolling the die.
+  """
+  result = random.randint(1, sides)
+  if "rolls" not in tool_context.state:
+    tool_context.state["rolls"] = []
+
+  tool_context.state["rolls"] = tool_context.state["rolls"] + [result]
+  return result
+
+
+vertex_ai_search_agent = Agent(
+    model="gemini-3-flash-preview",
+    name="vertex_ai_search_agent",
+    description="An agent for performing Vertex AI search.",
+    tools=[
+        VertexAiSearchTool(
+            data_store_id=VERTEXAI_DATASTORE_ID,
+        )
+    ],
+)
+
+root_agent = Agent(
+    model="gemini-3.1-pro-preview",
+    name="hello_world_agent",
+    description="A hello world agent with multiple tools.",
+    instruction="""
+      You are a helpful assistant which can help user to roll dice and search for information.
+      - Use `roll_die` tool to roll dice.
+      - Use `vertex_ai_search_agent` to search for Google Agent Development Kit (ADK) information in the datastore.
+    """,
+    tools=[
+        roll_die,
+        AgentTool(
+            agent=vertex_ai_search_agent, propagate_grounding_metadata=True
+        ),
+    ],
+)

src/google/adk/flows/llm_flows/base_llm_flow.py

@@ -259,7 +259,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

src/google/adk/tools/agent_tool.py

@@ -113,10 +113,12 @@ def __init__(
       skip_summarization: bool = False,
       *,
       include_plugins: bool = True,
+      propagate_grounding_metadata: bool = False,
   ):
     self.agent = agent
     self.skip_summarization: bool = skip_summarization
     self.include_plugins = include_plugins
+    self.propagate_grounding_metadata = propagate_grounding_metadata
 
     super().__init__(name=agent.name, description=agent.description)
 
@@ -247,6 +249,7 @@ async def run_async(
     )
 
     last_content = None
+    last_grounding_metadata = None
     async with Aclosing(
         runner.run_async(
             user_id=session.user_id, session_id=session.id, new_message=content
@@ -258,6 +261,7 @@ async def run_async(
           tool_context.state.update(event.actions.state_delta)
         if event.content:
           last_content = event.content
+          last_grounding_metadata = event.grounding_metadata
 
     # Clean up runner resources (especially MCP sessions)
     # to avoid "Attempted to exit cancel scope in a different task" errors
@@ -273,6 +277,12 @@ async def run_async(
       tool_result = validate_schema(output_schema, merged_text)
     else:
       tool_result = merged_text
+
+    if self.propagate_grounding_metadata and last_grounding_metadata:
+      tool_context.state['temp:_adk_grounding_metadata'] = (
+          last_grounding_metadata
+      )
+
     return tool_result
 
   @override

src/google/adk/tools/google_search_agent_tool.py

@@ -14,21 +14,12 @@
 
 from __future__ import annotations
 
-from typing import Any
 from typing import Union
 
-from google.genai import types
-from typing_extensions import override
-
 from ..agents.llm_agent import LlmAgent
-from ..memory.in_memory_memory_service import InMemoryMemoryService
 from ..models.base_llm import BaseLlm
-from ..utils._schema_utils import validate_schema
-from ..utils.context_utils import Aclosing
-from ._forwarding_artifact_service import ForwardingArtifactService
 from .agent_tool import AgentTool
 from .google_search_tool import google_search
-from .tool_context import ToolContext
 
 
 def create_google_search_agent(model: Union[str, BaseLlm]) -> LlmAgent:
@@ -60,80 +51,4 @@ class GoogleSearchAgentTool(AgentTool):
 
   def __init__(self, agent: LlmAgent):
     self.agent = agent
-    super().__init__(agent=self.agent)
-
-  @override
-  async def run_async(
-      self,
-      *,
-      args: dict[str, Any],
-      tool_context: ToolContext,
-  ) -> Any:
-    from ..agents.llm_agent import LlmAgent
-    from ..runners import Runner
-    from ..sessions.in_memory_session_service import InMemorySessionService
-
-    if isinstance(self.agent, LlmAgent) and self.agent.input_schema:
-      input_value = self.agent.input_schema.model_validate(args)
-      content = types.Content(
-          role='user',
-          parts=[
-              types.Part.from_text(
-                  text=input_value.model_dump_json(exclude_none=True)
-              )
-          ],
-      )
-    else:
-      content = types.Content(
-          role='user',
-          parts=[types.Part.from_text(text=args['request'])],
-      )
-    runner = Runner(
-        app_name=self.agent.name,
-        agent=self.agent,
-        artifact_service=ForwardingArtifactService(tool_context),
-        session_service=InMemorySessionService(),
-        memory_service=InMemoryMemoryService(),
-        credential_service=tool_context._invocation_context.credential_service,
-        plugins=list(tool_context._invocation_context.plugin_manager.plugins),
-    )
-
-    state_dict = {
-        k: v
-        for k, v in tool_context.state.to_dict().items()
-        if not k.startswith('_adk')  # Filter out adk internal states
-    }
-    session = await runner.session_service.create_session(
-        app_name=self.agent.name,
-        user_id=tool_context._invocation_context.user_id,
-        state=state_dict,
-    )
-
-    last_content = None
-    last_grounding_metadata = None
-    async with Aclosing(
-        runner.run_async(
-            user_id=session.user_id, session_id=session.id, new_message=content
-        )
-    ) as agen:
-      async for event in agen:
-        # Forward state delta to parent session.
-        if event.actions.state_delta:
-          tool_context.state.update(event.actions.state_delta)
-        if event.content:
-          last_content = event.content
-          last_grounding_metadata = event.grounding_metadata
-
-    if last_content is None or last_content.parts is None:
-      return ''
-    merged_text = '\n'.join(p.text for p in last_content.parts if p.text)
-    if isinstance(self.agent, LlmAgent) and self.agent.output_schema:
-      tool_result = validate_schema(self.agent.output_schema, merged_text)
-    else:
-      tool_result = merged_text
-
-    if last_grounding_metadata:
-      tool_context.state['temp:_adk_grounding_metadata'] = (
-          last_grounding_metadata
-      )
-    return tool_result
+    super().__init__(agent=self.agent, propagate_grounding_metadata=True)

tests/unittests/flows/llm_flows/test_base_llm_flow.py

@@ -284,7 +284,6 @@ async def test_handle_after_model_callback_grounding_with_no_callbacks(
       invocation_id=invocation_context.invocation_id,
       author=agent.name,
   )
-  flow = BaseLlmFlowForTesting()
 
   result = await _handle_after_model_callback(
       invocation_context, llm_response, event
@@ -341,7 +340,6 @@ async def test_handle_after_model_callback_grounding_with_callback_override(
       invocation_id=invocation_context.invocation_id,
       author=agent.name,
   )
-  flow = BaseLlmFlowForTesting()
 
   result = await _handle_after_model_callback(
       invocation_context, llm_response, event
@@ -403,7 +401,6 @@ def __init__(self):
       invocation_id=invocation_context.invocation_id,
       author=agent.name,
   )
-  flow = BaseLlmFlowForTesting()
 
   result = await _handle_after_model_callback(
       invocation_context, llm_response, event
@@ -430,6 +427,7 @@ class MockGoogleSearchTool(BaseTool):
 
       def __init__(self):
         super().__init__(name='google_search_agent', description='Mock search')
+        self.propagate_grounding_metadata = True
 
       async def call(self, **kwargs):
         return 'mock result'
@@ -459,7 +457,6 @@ async def call(self, **kwargs):
         invocation_id=invocation_context.invocation_id,
         author=agent.name,
     )
-    flow = BaseLlmFlowForTesting()
 
     # Call _handle_after_model_callback multiple times with the same context
     result1 = await _handle_after_model_callback(