feat: Support MCP prompts

Add support for MCP prompts via the McpInstructionProvider class, which can be specified as an agent's instruction. Co-authored-by: Kathy Wu <wukathy@google.com> PiperOrigin-RevId: 828166051

Author: wukath

contributing/samples/mcp_sse_agent/agent.py

@@ -16,25 +16,27 @@
 import os
 
 from google.adk.agents.llm_agent import LlmAgent
+from google.adk.agents.mcp_instruction_provider import McpInstructionProvider
 from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams
 from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset
 
 _allowed_path = os.path.dirname(os.path.abspath(__file__))
 
+connection_params = SseConnectionParams(
+ url='http://localhost:3000/sse',
+ headers={'Accept': 'text/event-stream'},
+)
+
 root_agent = LlmAgent(
  model='gemini-2.0-flash',
  name='enterprise_assistant',
- instruction=f"""\
-Help user accessing their file systems.
-
-Allowed directory: {_allowed_path}
- """,
+ instruction=McpInstructionProvider(
+ connection_params=connection_params,
+ prompt_name='file_system_prompt',
+ ),
  tools=[
  MCPToolset(
- connection_params=SseConnectionParams(
- url='http://localhost:3000/sse',
- headers={'Accept': 'text/event-stream'},
- ),
+ connection_params=connection_params,
  # don't want agent to do write operation
  # you can also do below
  # tool_filter=lambda tool, ctx=None: tool.name

contributing/samples/mcp_sse_agent/filesystem_server.py

@@ -45,6 +45,13 @@ def get_cwd() -> str:
  return str(Path.cwd())
 
 
+# Add a prompt for accessing file systems
+@mcp.prompt(name="file_system_prompt")
+def file_system_prompt() -> str:
+ return f"""\
+Help the user access their file systems."""
+
+
 # Graceful shutdown handler
 async def shutdown(signal, loop):
  """Cleanup tasks tied to the service's shutdown."""

src/google/adk/agents/__init__.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
+import sys
+
 from .base_agent import BaseAgent
 from .invocation_context import InvocationContext
 from .live_request_queue import LiveRequest
@@ -35,3 +38,16 @@
  'LiveRequestQueue',
  'RunConfig',
 ]
+
+if sys.version_info < (3, 10):
+ logger = logging.getLogger('google_adk.' + __name__)
+ logger.warning(
+ 'MCP requires Python 3.10 or above. Please upgrade your Python'
+ ' version in order to use it.'
+ )
+else:
+ from .mcp_instruction_provider import McpInstructionProvider
+
+ __all__.extend([
+ 'McpInstructionProvider',
+ ])

src/google/adk/agents/mcp_instruction_provider.py

@@ -0,0 +1,107 @@
+# Copyright 2025 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.
+
+"""Provides instructions to an agent by fetching prompts from an MCP server."""
+
+from __future__ import annotations
+
+import logging
+import sys
+from typing import Any
+from typing import Dict
+from typing import TextIO
+
+from .llm_agent import InstructionProvider
+from .readonly_context import ReadonlyContext
+
+# Attempt to import MCP Session Manager from the MCP library, and hints user to
+# upgrade their Python version to 3.10 if it fails.
+try:
+ from mcp import types
+
+ from ..tools.mcp_tool.mcp_session_manager import MCPSessionManager
+except ImportError as e:
+ if sys.version_info < (3, 10):
+ raise ImportError(
+ "MCP Session Manager requires Python 3.10 or above. Please upgrade"
+ " your Python version."
+ ) from e
+ else:
+ raise e
+
+
+class McpInstructionProvider(InstructionProvider):
+ """Fetches agent instructions from an MCP server."""
+
+ def __init__(
+ self,
+ connection_params: Any,
+ prompt_name: str,
+ errlog: TextIO = sys.stderr,
+ ):
+ """Initializes the McpInstructionProvider.
+
+ Args:
+ connection_params: Parameters for connecting to the MCP server.
+ prompt_name: The name of the MCP Prompt to fetch.
+ errlog: TextIO stream for error logging.
+ """
+ self._connection_params = connection_params
+ self._errlog = errlog or logging.getLogger(__name__)
+ self._mcp_session_manager = MCPSessionManager(
+ connection_params=self._connection_params,
+ errlog=self._errlog,
+ )
+ self.prompt_name = prompt_name
+
+ async def __call__(self, context: ReadonlyContext) -> str:
+ """Fetches the instruction from the MCP server.
+
+ Args:
+ context: The read-only context of the agent.
+
+ Returns:
+ The instruction string.
+ """
+ session = await self._mcp_session_manager.create_session()
+ # Fetch prompt definition to get the required argument names
+ prompt_definitions = await session.list_prompts()
+ prompt_definition = next(
+ (p for p in prompt_definitions.prompts if p.name == self.prompt_name),
+ None,
+ )
+
+ # Fetch arguments from context state if the prompt requires them
+ prompt_args: Dict[str, Any] = {}
+ if prompt_definition and prompt_definition.arguments:
+ arg_names = {arg.name for arg in prompt_definition.arguments}
+ prompt_args = {
+ k: v for k, v in (context.state or {}).items() if k in arg_names
+ }
+
+ # Fetch the specific prompt by name with arguments from context state
+ prompt_result: types.GetPromptResult = await session.get_prompt(
+ self.prompt_name, arguments=prompt_args
+ )
+
+ if prompt_result and prompt_result.messages:
+ # Concatenate content of all messages to form the instruction.
+ instruction = "".join(
+ message.content.text
+ for message in prompt_result.messages
+ if message.content.type == "text"
+ )
+ return instruction
+ else:
+ raise ValueError(f"Failed to load MCP prompt '{self.prompt_name}'.")

tests/unittests/agents/test_mcp_instruction_provider.py

@@ -0,0 +1,210 @@
+# Copyright 2025 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.
+
+"""Unit tests for McpInstructionProvider."""
+import sys
+from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
+from unittest.mock import patch
+
+from google.adk.agents.readonly_context import ReadonlyContext
+import pytest
+
+# Skip all tests in this module if Python version is less than 3.10
+pytestmark = pytest.mark.skipif(
+ sys.version_info < (3, 10),
+ reason="MCP instruction provider requires Python 3.10+",
+)
+
+# Import dependencies with version checking
+try:
+ from google.adk.agents.mcp_instruction_provider import McpInstructionProvider
+except ImportError as e:
+ if sys.version_info < (3, 10):
+ # Create dummy classes to prevent NameError during test collection
+ # Tests will be skipped anyway due to pytestmark
+ class DummyClass:
+ pass
+
+ McpInstructionProvider = DummyClass
+ else:
+ raise e
+
+
+class TestMcpInstructionProvider:
+ """Unit tests for McpInstructionProvider."""
+
+ def setup_method(self):
+ """Sets up the test environment."""
+ self.connection_params = {"host": "localhost", "port": 8000}
+ self.prompt_name = "test_prompt"
+ self.mock_mcp_session_manager_cls = patch(
+ "google.adk.agents.mcp_instruction_provider.MCPSessionManager"
+ ).start()
+ self.mock_mcp_session_manager = (
+ self.mock_mcp_session_manager_cls.return_value
+ )
+ self.mock_session = MagicMock()
+ self.mock_session.list_prompts = AsyncMock()
+ self.mock_session.get_prompt = AsyncMock()
+ self.mock_mcp_session_manager.create_session = AsyncMock(
+ return_value=self.mock_session
+ )
+ self.provider = McpInstructionProvider(
+ self.connection_params, self.prompt_name
+ )
+
+ @pytest.mark.asyncio
+ async def test_call_success_no_args(self):
+ """Tests __call__ with a prompt that has no arguments."""
+ mock_prompt = MagicMock()
+ mock_prompt.name = self.prompt_name
+ mock_prompt.arguments = None
+ self.mock_session.list_prompts.return_value = MagicMock(
+ prompts=[mock_prompt]
+ )
+
+ mock_msg1 = MagicMock()
+ mock_msg1.content.type = "text"
+ mock_msg1.content.text = "instruction part 1. "
+ mock_msg2 = MagicMock()
+ mock_msg2.content.type = "text"
+ mock_msg2.content.text = "instruction part 2"
+ self.mock_session.get_prompt.return_value = MagicMock(
+ messages=[mock_msg1, mock_msg2]
+ )
+
+ mock_invocation_context = MagicMock()
+ mock_invocation_context.session.state = {}
+ context = ReadonlyContext(mock_invocation_context)
+
+ # Call
+ instruction = await self.provider(context)
+
+ # Assert
+ assert instruction == "instruction part 1. instruction part 2"
+ self.mock_session.get_prompt.assert_called_once_with(
+ self.prompt_name, arguments={}
+ )
+
+ @pytest.mark.asyncio
+ async def test_call_success_with_args(self):
+ """Tests __call__ with a prompt that has arguments."""
+ mock_arg1 = MagicMock()
+ mock_arg1.name = "arg1"
+ mock_prompt = MagicMock()
+ mock_prompt.name = self.prompt_name
+ mock_prompt.arguments = [mock_arg1]
+ self.mock_session.list_prompts.return_value = MagicMock(
+ prompts=[mock_prompt]
+ )
+
+ mock_msg = MagicMock()
+ mock_msg.content.type = "text"
+ mock_msg.content.text = "instruction with arg1"
+ self.mock_session.get_prompt.return_value = MagicMock(messages=[mock_msg])
+
+ mock_invocation_context = MagicMock()
+ mock_invocation_context.session.state = {"arg1": "value1", "arg2": "value2"}
+ context = ReadonlyContext(mock_invocation_context)
+
+ instruction = await self.provider(context)
+
+ assert instruction == "instruction with arg1"
+ self.mock_session.get_prompt.assert_called_once_with(
+ self.prompt_name, arguments={"arg1": "value1"}
+ )
+
+ @pytest.mark.asyncio
+ async def test_call_prompt_not_found_in_list_prompts(self):
+ """Tests __call__ when list_prompts doesn't return the prompt."""
+ self.mock_session.list_prompts.return_value = MagicMock(prompts=[])
+
+ mock_msg = MagicMock()
+ mock_msg.content.type = "text"
+ mock_msg.content.text = "instruction"
+ self.mock_session.get_prompt.return_value = MagicMock(messages=[mock_msg])
+
+ mock_invocation_context = MagicMock()
+ mock_invocation_context.session.state = {"arg1": "value1"}
+ context = ReadonlyContext(mock_invocation_context)
+
+ instruction = await self.provider(context)
+
+ assert instruction == "instruction"
+ self.mock_session.get_prompt.assert_called_once_with(
+ self.prompt_name, arguments={}
+ )
+
+ @pytest.mark.asyncio
+ async def test_call_get_prompt_returns_no_messages(self):
+ """Tests __call__ when get_prompt returns no messages."""
+ # Setup mocks
+ self.mock_session.list_prompts.return_value = MagicMock(prompts=[])
+ self.mock_session.get_prompt.return_value = MagicMock(messages=[])
+
+ mock_invocation_context = MagicMock()
+ mock_invocation_context.session.state = {}
+ context = ReadonlyContext(mock_invocation_context)
+
+ # Call and assert
+ with pytest.raises(
+ ValueError, match="Failed to load MCP prompt 'test_prompt'."
+ ):
+ await self.provider(context)
+
+ # Assert
+ self.mock_session.get_prompt.assert_called_once_with(
+ self.prompt_name, arguments={}
+ )
+
+ @pytest.mark.asyncio
+ async def test_call_ignore_non_text_messages(self):
+ """Tests __call__ ignores non-text messages."""
+ # Setup mocks
+ mock_prompt = MagicMock()
+ mock_prompt.name = self.prompt_name
+ mock_prompt.arguments = None
+ self.mock_session.list_prompts.return_value = MagicMock(
+ prompts=[mock_prompt]
+ )
+
+ mock_msg1 = MagicMock()
+ mock_msg1.content.type = "text"
+ mock_msg1.content.text = "instruction part 1. "
+
+ mock_msg2 = MagicMock()
+ mock_msg2.content.type = "image"
+ mock_msg2.content.text = "ignored"
+
+ mock_msg3 = MagicMock()
+ mock_msg3.content.type = "text"
+ mock_msg3.content.text = "instruction part 2"
+
+ self.mock_session.get_prompt.return_value = MagicMock(
+ messages=[mock_msg1, mock_msg2, mock_msg3]
+ )
+
+ mock_invocation_context = MagicMock()
+ mock_invocation_context.session.state = {}
+ context = ReadonlyContext(mock_invocation_context)
+
+ # Call
+ instruction = await self.provider(context)
+
+ # Assert
+ assert instruction == "instruction part 1. instruction part 2"
+ self.mock_session.get_prompt.assert_called_once_with(
+ self.prompt_name, arguments={}
+ )