[Feature Request] Support grounding metadata propagation to root agent when routing between multiple VertexAiSearchTool instances

[NNNNN11111111]

This issue is a follow-up to the discussion and findings in Issue #3406 (cc: @xuanyang15).

🔴 Required Information

Is your feature request related to a specific problem?

Yes. I am trying to use multiple VertexAiSearchTool instances (e.g., one for structured data, one for unstructured data) in a single agent so the LLM can dynamically route user queries to the appropriate datastore.
Because the built-in tool's name is hardcoded to vertex_ai_search, the LLM treats all instances as a single tool and cannot route queries. When I try the workaround of wrapping them in sub-agents via a custom AgentTool (to give them distinct names and descriptions for routing), the routing works perfectly, but the grounding metadata (citations) is lost when returning to the root agent.

Describe the Solution You'd Like

I would like the root agent to be able to successfully merge and retain grounding metadata propagated from a sub-agent, so that citations are correctly displayed in the final response (e.g., in the ADK Web UI).
Alternatively, providing an official VertexAiSearchAgentTool that handles this metadata propagation out-of-the-box (similar to GoogleSearchAgentTool) would perfectly solve this.

Impact on your work

This is a major blocker for our use case. Our goal is to build a single smart agent that can route queries to the correct Vertex AI Datastore while strictly providing grounded answers with citations. Without metadata propagation from sub-agents, we are forced to choose between "intelligent routing" and "retaining citations," as we cannot achieve both simultaneously.

Willingness to contribute

No

---

🟡 Recommended Information

Describe Alternatives You've Considered

• Attempt 1: All tools in the root agent. * Why it failed: Built-in tools do not accept custom names. The LLM cannot distinguish between the datastores, making dynamic routing impossible.
• Attempt 2: Custom AgentTool propagating metadata (Suggested in Issue with structured VertexAI datastores when using multiple instances of VertexAiSearchTool in the same agent #3406). * Why it failed: Following @xuanyang15's advice, I created a custom MetadataPropagatingAgentTool based on GoogleSearchAgentTool, adding tool_context.state['temp:_adk_grounding_metadata'] = last_grounding_metadata. While local debug logs confirmed the sub-agent successfully queries the right datastore and retrieves the metadata, the root agent fails to merge this metadata into the final response. The citations are missing in the UI, and the root LLM lacks the context of the retrieved documents (it cannot answer "Which document did you refer to?").

Proposed API / Implementation

Update the root agent's response generation logic (or artifact service) to respect and merge grounding metadata returned from sub-agent tool executions, ensuring it gets passed down to the final UI response.

Additional Context

• ADK Version: v1.25.0
• Model: gemini-3-flash-preview

[NNNNN11111111]

To help resolve this issue, I have drafted a potential solution using a custom AgentTool and a custom Plugin.

I was able to successfully propagate the grounding_metadata from a VertexAiSearchTool (running inside a sub-agent) back to the root agent's response. I am sharing my implementation here as a reference or a possible starting point for an official fix.

1) Custom AgentTool (Saves metadata to state)

We use a custom AgentTool to extract the grounding_metadata during the sub-agent's execution and temporarily store it in tool_context.state, keyed by the function_call_id.

# grounding_agent_tool.py
from __future__ import annotations
from typing import Any

from google.adk.agents.llm_agent import LlmAgent
from google.adk.memory.in_memory_memory_service import InMemoryMemoryService
from google.adk.runners import Runner
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.adk.tools._forwarding_artifact_service import ForwardingArtifactService
from google.adk.tools.agent_tool import AgentTool
from google.adk.tools.tool_context import ToolContext
from google.adk.utils.context_utils import Aclosing
from google.genai import types
from typing_extensions import override

GROUNDING_CALL_STATE_KEY_PREFIX = "temp:grounding_bridge:call:"

class GroundingBridgeAgentTool(AgentTool):
    """An AgentTool that saves the sub-agent's grounding metadata into the state."""

    def __init__(
        self,
        agent: LlmAgent,
        *,
        name: str | None = None,
        description: str | None = None,
        state_key_prefix: str = GROUNDING_CALL_STATE_KEY_PREFIX,
        include_plugins: bool = True,
        skip_summarization: bool = False,
    ):
        self.agent = agent
        self._state_key_prefix = state_key_prefix
        super().__init__(
            agent=self.agent,
            include_plugins=include_plugins,
            skip_summarization=skip_summarization,
        )
        if name:
            self.name = name
        if description:
            self.description = description

    @override
    async def run_async(
        self,
        *,
        args: dict[str, Any],
        tool_context: ToolContext,
    ) -> Any:
        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),
        )

        # Do not pass ADK internal keys (_adk*) to the sub-session
        state_dict = {
            k: v
            for k, v in tool_context.state.to_dict().items()
            if not k.startswith("_adk")
        }
        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: dict[str, Any] | Any | None = 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:
                if event.actions.state_delta:
                    tool_context.state.update(event.actions.state_delta)

                if event.grounding_metadata:
                    if hasattr(event.grounding_metadata, "model_dump"):
                        last_grounding_metadata = event.grounding_metadata.model_dump(exclude_none=True)
                    else:
                        last_grounding_metadata = event.grounding_metadata

                if event.content:
                    last_content = event.content

        await runner.close()

        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 and not p.thought
        )

        if isinstance(self.agent, LlmAgent) and self.agent.output_schema:
            tool_result = self.agent.output_schema.model_validate_json(merged_text).model_dump(exclude_none=True)
        else:
            tool_result = merged_text

        function_call_id = tool_context.function_call_id
        if last_grounding_metadata and function_call_id:
            state_key = f"{self._state_key_prefix}{function_call_id}"
            tool_context.state[state_key] = {
                "tool_name": self.name,
                "agent_name": self.agent.name,
                "function_call_id": function_call_id,
                "grounding_metadata": last_grounding_metadata,
            }

        return tool_result

2) Plugin (Injects metadata back into the root response)

Using after_tool_callback, we retrieve the saved metadata from the state and inject it directly into the tool's response payload (result), ensuring the root agent receives the citations.

# grounding_plugin.py
from __future__ import annotations
from datetime import datetime, timezone
from typing import Any

from google.adk.plugins.base_plugin import BasePlugin
from google.adk.tools.base_tool import BaseTool
from google.adk.tools.tool_context import ToolContext
from .grounding_agent_tool import GROUNDING_CALL_STATE_KEY_PREFIX

DEFAULT_GROUNDING_AGGREGATE_STATE_KEY = "temp:grounding_bridge:aggregate"
DEFAULT_GROUNDING_HISTORY_STATE_KEY = "temp:grounding_bridge:history"

def _as_dict(value: Any) -> dict[str, Any]:
    if isinstance(value, dict):
        return value
    if hasattr(value, "model_dump"):
        return value.model_dump(exclude_none=True)
    return {}

def _as_list(value: Any) -> list[Any]:
    return value if isinstance(value, list) else []

def merge_references(current: list[dict[str, Any]], new_items: list[dict[str, Any]]) -> list[dict[str, Any]]:
    merged: list[dict[str, Any]] = []
    seen: set[tuple[str, str, str, str]] = set()
    for ref in [*current, *new_items]:
        key = (
            str(ref.get("uri") or ""),
            str(ref.get("document_name") or ""),
            str(ref.get("title") or ""),
            str(ref.get("source_tool") or ""),
        )
        if key in seen:
            continue
        seen.add(key)
        merged.append(ref)
    return merged

def build_references(grounding_metadata: dict[str, Any], *, source_tool: str) -> list[dict[str, Any]]:
    refs: list[dict[str, Any]] = []
    chunks = _as_list(grounding_metadata.get("grounding_chunks"))
    for idx, chunk in enumerate(chunks):
        c = _as_dict(chunk)
        rc = _as_dict(c.get("retrieved_context"))
        if not rc:
            continue
        refs.append(
            {
                "source_tool": source_tool,
                "title": rc.get("title") or rc.get("document_name") or f"Reference {idx + 1}",
                "uri": rc.get("uri") or "",
                "url": rc.get("uri") or "",
                "document_name": rc.get("document_name") or "",
                "chunk_index": idx,
            }
        )
    return merge_references([], refs)

class GroundingBridgePlugin(BasePlugin):
    """Injects grounding metadata saved by the AgentTool into the tool_response."""

    def __init__(
        self,
        *,
        name: str = "grounding_bridge_plugin",
        state_key_prefix: str = GROUNDING_CALL_STATE_KEY_PREFIX,
        aggregate_state_key: str = DEFAULT_GROUNDING_AGGREGATE_STATE_KEY,
        history_state_key: str = DEFAULT_GROUNDING_HISTORY_STATE_KEY,
    ):
        super().__init__(name=name)
        self._state_key_prefix = state_key_prefix
        self._aggregate_state_key = aggregate_state_key
        self._history_state_key = history_state_key

    async def after_tool_callback(
        self,
        *,
        tool: BaseTool,
        tool_args: dict[str, Any],
        tool_context: ToolContext,
        result: dict[str, Any],
    ) -> dict[str, Any] | None:
        function_call_id = tool_context.function_call_id
        if not function_call_id:
            return None

        state_key = f"{self._state_key_prefix}{function_call_id}"
        call_record = _as_dict(tool_context.state.get(state_key))
        grounding_metadata = _as_dict(call_record.get("grounding_metadata"))
        if not grounding_metadata:
            return None

        source_tool = str(call_record.get("tool_name") or tool.name)
        references = build_references(grounding_metadata, source_tool=source_tool)

        response = dict(result) if isinstance(result, dict) else {"result": result}
        response["grounding_metadata"] = grounding_metadata
        response["references"] = references

        history = _as_list(tool_context.state.get(self._history_state_key))
        history.append(
            {
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "tool_name": source_tool,
                "function_call_id": function_call_id,
                "args": tool_args,
                "references": references,
            }
        )
        tool_context.state[self._history_state_key] = history

        aggregate = _as_dict(tool_context.state.get(self._aggregate_state_key))
        current_refs = _as_list(aggregate.get("references"))
        merged_refs = merge_references(current_refs, references)
        tool_context.state[self._aggregate_state_key] = {
            "updated_at": datetime.now(timezone.utc).isoformat(),
            "reference_count": len(merged_refs),
            "references": merged_refs,
        }

        return response

3) Usage (Minimal Example)

tool_abc = GroundingBridgeAgentTool(agent=sub_agent_abc, name="search_abc_datastore")
tool_xyz = GroundingBridgeAgentTool(agent=sub_agent_xyz, name="search_xyz_datastore")

root_agent = LlmAgent(..., tools=[tool_abc, tool_xyz])

app = App(
    name="grounding_bridge_demo",
    root_agent=root_agent,
    plugins=[GroundingBridgePlugin()],
)

I have confirmed that this method correctly enables the root LLM response to handle and display grounding_metadata and references propagated from the sub-agents via the AgentTool.

[klateefa]

@NNNNN11111111 Thank you for submitting this feature request. It’s currently under review by our team, and I’ll keep you posted on any updates. We appreciate your input!

[xuanyang15]

@NNNNN11111111 Thanks for reporting this issue and the investigation! I have confirmed with the model team that it is indeed a limitation with the current model, even if the model API can add a name/id for the two datastores, the model still cannot dynamically choose which one to use because of how the grounding metadata is currently being injected for the model.

So, it looks we have to workaround on ADK level, I will take a closer look at you draft solution and also talk with ADK team to come up with a final solution.