feat: Add support for MCP App UI widgets in MCPTool

Co-authored-by: Vishwa Murugan <vishwamurugan@google.com>
PiperOrigin-RevId: 881131683

Author: vishwa180

src/google/adk/agents/context.py

@@ -31,6 +31,7 @@
   from ..auth.auth_tool import AuthConfig
   from ..events.event import Event
   from ..events.event_actions import EventActions
+  from ..events.ui_widget import UiWidget
   from ..memory.base_memory_service import SearchMemoryResponse
   from ..memory.memory_entry import MemoryEntry
   from ..sessions.state import State
@@ -410,3 +411,29 @@ async def search_memory(self, query: str) -> SearchMemoryResponse:
         user_id=self._invocation_context.user_id,
         query=query,
     )
+
+  # ============================================================================
+  # UI Widget methods
+  # ============================================================================
+
+  def render_ui_widget(self, ui_widget: UiWidget) -> None:
+    """Adds a UI widget to the current event's actions for the UI to render.
+
+    UI widgets provide rendering payload/metadata that the UI Host uses to
+    display rich interactive components (e.g., MCP App iframes) alongside agent
+    responses.
+
+    Args:
+      ui_widget: A ``UiWidget`` instance.
+    """
+    if self._event_actions.render_ui_widgets is None:
+      self._event_actions.render_ui_widgets = []
+
+    for existing_widget in self._event_actions.render_ui_widgets:
+      if existing_widget.id == ui_widget.id:
+        raise ValueError(
+            f"UI widget with ID '{ui_widget.id}' already exists in the current"
+            " event actions."
+        )
+
+    self._event_actions.render_ui_widgets.append(ui_widget)

src/google/adk/tools/mcp_tool/mcp_tool.py

@@ -21,6 +21,7 @@
 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
@@ -39,6 +40,7 @@
 from ...auth.auth_credential import AuthCredential
 from ...auth.auth_schemes import AuthScheme
 from ...auth.auth_tool import AuthConfig
+from ...events.ui_widget import UiWidget
 from ...features import FeatureName
 from ...features import is_feature_enabled
 from ...utils.context_utils import find_context_parameter
@@ -212,6 +214,48 @@ def raw_mcp_tool(self) -> McpBaseTool:
     """Returns the raw MCP tool."""
     return self._mcp_tool
 
+  @property
+  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):
+      return []
+
+    # Format: meta.ui.visibility
+    ui = meta.get("ui", {})
+    if isinstance(ui, dict):
+      return ui.get("visibility", [])
+    return []
+
+  @property
+  def mcp_app_resource_uri(self) -> Optional[str]:
+    """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
+    definition. This property extracts that URI, supporting both the nested
+    format (`{"ui": {"resourceUri": "ui://..."}}`) and the flat format
+    (`{"ui/resourceUri": "ui://..."}`).
+
+    Returns:
+        The `ui://` resource URI string, or None if not present.
+    """
+    meta = getattr(self.raw_mcp_tool, "meta", None)
+    if not meta or not isinstance(meta, dict):
+      return None
+    # Nested format: meta.ui.resourceUri (preferred)
+    ui = meta.get("ui")
+    if isinstance(ui, dict):
+      uri = ui.get("resourceUri")
+      if isinstance(uri, str) and uri.startswith("ui://"):
+        return uri
+    # Flat format: meta["ui/resourceUri"] (deprecated)
+    # Reference:
+    # https://github.com/modelcontextprotocol/ext-apps/blob/main/specification/2026-01-26/apps.mdx
+    uri = meta.get("ui/resourceUri")
+    if isinstance(uri, str) and uri.startswith("ui://"):
+      return uri
+    return None
+
   async def _invoke_callable(
       self, target: Callable[..., Any], args_to_call: dict[str, Any]
   ) -> Any:
@@ -335,7 +379,22 @@ async def _run_async_impl(
         progress_callback=resolved_callback,
         meta=meta_trace_context,
     )
-    return response.model_dump(exclude_none=True, mode="json")
+    result = response.model_dump(exclude_none=True, mode="json")
+
+    # Push UI widget to the event actions if the tool supports it.
+    if self.mcp_app_resource_uri:
+      tool_context.render_ui_widget(
+          UiWidget(
+              id=tool_context.function_call_id,
+              provider="mcp",
+              payload={
+                  "resource_uri": self.mcp_app_resource_uri,
+                  "tool": self._mcp_tool,
+                  "tool_args": args,
+              },
+          )
+      )
+    return result
 
   def _resolve_progress_callback(
       self, tool_context: ToolContext

tests/unittests/agents/test_context.py

@@ -21,6 +21,7 @@
 from google.adk.auth.auth_credential import AuthCredential
 from google.adk.auth.auth_credential import AuthCredentialTypes
 from google.adk.auth.auth_tool import AuthConfig
+from google.adk.events.ui_widget import UiWidget
 from google.adk.memory.base_memory_service import SearchMemoryResponse
 from google.adk.memory.memory_entry import MemoryEntry
 from google.adk.tools.tool_confirmation import ToolConfirmation
@@ -547,3 +548,74 @@ async def test_add_memory_no_service_raises(self, mock_invocation_context):
               )
           ]
       )
+
+
+class TestContextAddUiWidget:
+  """Test render_ui_widget method in Context."""
+
+  def test_render_ui_widget(self, mock_invocation_context):
+    """Test that render_ui_widget appends a widget to actions."""
+
+    context = Context(mock_invocation_context)
+    widget = UiWidget(
+        id="w1",
+        provider="mcp",
+        payload={"resource_uri": "ui://test-app"},
+    )
+
+    context.render_ui_widget(widget)
+
+    assert context.actions.render_ui_widgets is not None
+    assert len(context.actions.render_ui_widgets) == 1
+    assert context.actions.render_ui_widgets[0] is widget
+
+  def test_render_ui_widget_multiple(self, mock_invocation_context):
+    """Test that calling render_ui_widget twice yields two widgets."""
+
+    context = Context(mock_invocation_context)
+    w1 = UiWidget(
+        id="w1",
+        provider="mcp",
+        payload={"resource_uri": "ui://app-1"},
+    )
+    w2 = UiWidget(
+        id="w2",
+        provider="mcp",
+        payload={"resource_uri": "ui://app-2"},
+    )
+
+    context.render_ui_widget(w1)
+    context.render_ui_widget(w2)
+
+    assert len(context.actions.render_ui_widgets) == 2
+    assert context.actions.render_ui_widgets[0] is w1
+    assert context.actions.render_ui_widgets[1] is w2
+
+  def test_render_ui_widget_duplicate(self, mock_invocation_context):
+    """Test that duplicate widgets by id are not added."""
+
+    context = Context(mock_invocation_context)
+    w1 = UiWidget(
+        id="w1",
+        provider="mcp",
+        payload={"resource_uri": "ui://app-1"},
+    )
+    w2 = UiWidget(
+        id="w1",
+        provider="mcp",
+        payload={"resource_uri": "ui://app-1-mod"},
+    )
+
+    context.render_ui_widget(w1)
+
+    with pytest.raises(
+        ValueError,
+        match=(
+            f"UI widget with ID '{w1.id}' already exists in the current event"
+            " actions."
+        ),
+    ):
+      context.render_ui_widget(w2)
+
+    assert len(context.actions.render_ui_widgets) == 1
+    assert context.actions.render_ui_widgets[0] is w1

tests/unittests/tools/mcp_tool/test_mcp_tool.py

@@ -46,9 +46,11 @@ def __init__(
       name="test_tool",
       description="Test tool description",
       outputSchema=None,
+      meta=None,
   ):
     self.name = name
     self.description = description
+    self.meta = meta
     self.inputSchema = {
         "type": "object",
         "properties": {
@@ -213,7 +215,8 @@ async def test_run_async_impl_no_auth(self):
     )
     self.mock_session.call_tool = AsyncMock(return_value=mcp_response)
 
-    tool_context = Mock(spec=ToolContext)
+    tool_context = ToolContext(invocation_context=Mock())
+    tool_context.function_call_id = "test-call-id"
     args = {"param1": "test_value"}
 
     result = await tool._run_async_impl(
@@ -230,6 +233,42 @@ async def test_run_async_impl_no_auth(self):
         "test_tool", arguments=args, progress_callback=None, meta=None
     )
 
+  @pytest.mark.asyncio
+  async def test_run_async_impl_adds_ui_widget(self):
+    """Test running tool adds UiWidget to actions."""
+    meta = {"ui": {"resourceUri": "ui://test-app"}}
+    mock_tool = MockMCPTool(meta=meta)
+    tool = MCPTool(
+        mcp_tool=mock_tool,
+        mcp_session_manager=self.mock_session_manager,
+    )
+
+    mcp_response = CallToolResult(
+        content=[TextContent(type="text", text="success")]
+    )
+    self.mock_session.call_tool = AsyncMock(return_value=mcp_response)
+
+    tool_context = ToolContext(invocation_context=Mock())
+    tool_context.function_call_id = "test-call-id"
+    args = {"param1": "test_value"}
+
+    # tool_context.actions.render_ui_widgets is None initially
+    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")
+
+    assert tool_context.actions.render_ui_widgets is not None
+    assert len(tool_context.actions.render_ui_widgets) == 1
+    widget = tool_context.actions.render_ui_widgets[0]
+
+    assert widget.id == "test-call-id"
+    assert widget.provider == "mcp"
+    assert widget.payload["resource_uri"] == "ui://test-app"
+    assert widget.payload["tool"] == mock_tool
+    assert widget.payload["tool_args"] == args
+
   @pytest.mark.asyncio
   async def test_run_async_impl_with_oauth2(self):
     """Test running tool with OAuth2 authentication."""
@@ -1015,3 +1054,76 @@ async def _require_confirmation_func(param1: str, ctx: Context):
           )
       }
       tool_context.request_confirmation.assert_called_once()
+
+  def test_visibility_property(self):
+    """Test visibility property extraction from meta."""
+    meta = {"ui": {"visibility": ["app", "debug"]}}
+    mock_tool = MockMCPTool(meta=meta)
+    tool = MCPTool(
+        mcp_tool=mock_tool,
+        mcp_session_manager=self.mock_session_manager,
+    )
+
+    assert tool.visibility == ["app", "debug"]
+
+  def test_visibility_property_empty(self):
+    """Test visibility property when meta is missing or malformed."""
+    # Missing meta
+    tool1 = MCPTool(
+        mcp_tool=MockMCPTool(meta=None),
+        mcp_session_manager=self.mock_session_manager,
+    )
+    assert tool1.visibility == []
+
+    # Malformed meta
+    tool2 = MCPTool(
+        mcp_tool=MockMCPTool(meta="not a dict"),
+        mcp_session_manager=self.mock_session_manager,
+    )
+    assert tool2.visibility == []
+
+    # Missing ui field
+    tool3 = MCPTool(
+        mcp_tool=MockMCPTool(meta={}),
+        mcp_session_manager=self.mock_session_manager,
+    )
+    assert tool3.visibility == []
+
+  def test_mcp_app_resource_uri_property_nested(self):
+    """Test MCP App resource URI extraction from nested meta format."""
+    meta = {"ui": {"resourceUri": "ui://test-resource"}}
+    mock_tool = MockMCPTool(meta=meta)
+    tool = MCPTool(
+        mcp_tool=mock_tool,
+        mcp_session_manager=self.mock_session_manager,
+    )
+
+    assert tool.mcp_app_resource_uri == "ui://test-resource"
+
+  def test_mcp_app_resource_uri_property_flat(self):
+    """Test MCP App resource URI extraction from flat meta format."""
+    meta = {"ui/resourceUri": "ui://test-resource-flat"}
+    mock_tool = MockMCPTool(meta=meta)
+    tool = MCPTool(
+        mcp_tool=mock_tool,
+        mcp_session_manager=self.mock_session_manager,
+    )
+
+    assert tool.mcp_app_resource_uri == "ui://test-resource-flat"
+
+  def test_mcp_app_resource_uri_property_none(self):
+    """Test MCP App resource URI when missing or invalid."""
+    # Missing meta
+    tool1 = MCPTool(
+        mcp_tool=MockMCPTool(meta=None),
+        mcp_session_manager=self.mock_session_manager,
+    )
+    assert tool1.mcp_app_resource_uri is None
+
+    # Invalid scheme
+    meta = {"ui": {"resourceUri": "http://invalid"}}
+    tool2 = MCPTool(
+        mcp_tool=MockMCPTool(meta=meta),
+        mcp_session_manager=self.mock_session_manager,
+    )
+    assert tool2.mcp_app_resource_uri is None