mirror of
https://github.com/microsoft/semantic-kernel.git
synced 2026-10-12 02:19:06 +00:00
Python: Improvements for MCP (#14003)
### Motivation and Context Improvements for MCP tool calls by adding an optional sampling callback. <!-- Thank you for your contribution to the semantic-kernel repo! Please help reviewers and future users, providing the following information: 1. Why is this change required? 2. What problem does it solve? 3. What scenario does it contribute to? 4. If it fixes an open issue, please link to the issue here. --> ### Description Improvements for MCP tool calls by adding an optional sampling callback. <!-- Describe your changes, the overall approach, the underlying design. These notes will help understanding how your code works. Thanks! --> ### Contribution Checklist <!-- Before submitting this PR, please make sure: --> - [X] The code builds clean without any errors or warnings - [X] The PR follows the [SK Contribution Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md) and the [pre-submission formatting script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts) raises no violations - [X] All unit tests pass, and I have added new tests where possible - [ ] I didn't break anyone 😄 --------- Co-authored-by: Copilot <[email protected]>
This commit is contained in:
@@ -6,7 +6,7 @@ import logging
|
||||
import re
|
||||
import sys
|
||||
from abc import abstractmethod
|
||||
from collections.abc import Callable, Sequence
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from contextlib import AbstractAsyncContextManager, AsyncExitStack, _AsyncGeneratorContextManager
|
||||
from datetime import timedelta
|
||||
from functools import partial
|
||||
@@ -59,6 +59,8 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
# region: Helpers
|
||||
|
||||
SamplingConsentCallback = Callable[[str, types.CreateMessageRequestParams], Awaitable[bool]]
|
||||
|
||||
LOG_LEVEL_MAPPING: dict[types.LoggingLevel, int] = {
|
||||
"debug": logging.DEBUG,
|
||||
"info": logging.INFO,
|
||||
@@ -243,8 +245,23 @@ class MCPPluginBase:
|
||||
session: ClientSession | None = None,
|
||||
kernel: Kernel | None = None,
|
||||
request_timeout: int | None = None,
|
||||
sampling_consent_callback: SamplingConsentCallback | None = None,
|
||||
) -> None:
|
||||
"""Initialize the MCP Plugin Base."""
|
||||
"""Initialize the MCP Plugin Base.
|
||||
|
||||
Args:
|
||||
name: The name of the plugin.
|
||||
description: The description of the plugin.
|
||||
load_tools: Whether to load tools from the MCP server.
|
||||
load_prompts: Whether to load prompts from the MCP server.
|
||||
session: The session to use for the MCP connection.
|
||||
kernel: The kernel instance with one or more Chat Completion clients.
|
||||
request_timeout: The default timeout used for all requests.
|
||||
sampling_consent_callback: Optional callback for approving MCP sampling requests.
|
||||
Receives the plugin name and MCP sampling request params. Return
|
||||
False to deny the request. When omitted, sampling requests are
|
||||
auto-approved and a warning is logged.
|
||||
"""
|
||||
self.name = name
|
||||
self.description = description
|
||||
self.load_tools_flag = load_tools
|
||||
@@ -253,6 +270,9 @@ class MCPPluginBase:
|
||||
self.session = session
|
||||
self.kernel = kernel or None
|
||||
self.request_timeout = request_timeout
|
||||
self.sampling_consent_callback = sampling_consent_callback
|
||||
self._sampling_auto_approved_warning_logged = False
|
||||
self._mcp_reserved_attribute_names: set[str] | None = None
|
||||
self._current_task: asyncio.Task | None = None
|
||||
self._stop_event: asyncio.Event | None = None
|
||||
|
||||
@@ -361,9 +381,24 @@ class MCPPluginBase:
|
||||
|
||||
This function is called when the MCP server needs to get a message completed.
|
||||
|
||||
This is a simple version of this function, it can be overridden to allow more complex sampling.
|
||||
It get's added to the session at initialization time, so overriding it is the best way to do this.
|
||||
If a sampling consent callback is configured, it is called before forwarding the request to the configured
|
||||
chat completion service. Returning False denies the request. If no callback is configured, requests are
|
||||
auto-approved and a warning is logged.
|
||||
"""
|
||||
if self.sampling_consent_callback is None:
|
||||
if not self._sampling_auto_approved_warning_logged:
|
||||
logger.warning(
|
||||
"MCP sampling request for plugin '%s' was auto-approved because no sampling consent callback "
|
||||
"was configured.",
|
||||
self.name,
|
||||
)
|
||||
self._sampling_auto_approved_warning_logged = True
|
||||
elif not await self._is_sampling_approved(params):
|
||||
return types.ErrorData(
|
||||
code=types.INTERNAL_ERROR,
|
||||
message="Sampling denied by policy.",
|
||||
)
|
||||
|
||||
if not self.kernel or not self.kernel.services:
|
||||
return types.ErrorData(
|
||||
code=types.INTERNAL_ERROR,
|
||||
@@ -431,6 +466,15 @@ class MCPPluginBase:
|
||||
model=service.ai_model_id,
|
||||
)
|
||||
|
||||
async def _is_sampling_approved(self, params: types.CreateMessageRequestParams) -> bool:
|
||||
if self.sampling_consent_callback is None:
|
||||
return True
|
||||
try:
|
||||
return await self.sampling_consent_callback(self.name, params)
|
||||
except Exception:
|
||||
logger.exception("MCP sampling consent callback failed for plugin '%s'.", self.name)
|
||||
return False
|
||||
|
||||
async def logging_callback(self, params: types.LoggingMessageNotificationParams) -> None:
|
||||
"""Callback function for logging.
|
||||
|
||||
@@ -464,6 +508,19 @@ class MCPPluginBase:
|
||||
case "notifications/prompts/list_changed":
|
||||
await self.load_prompts()
|
||||
|
||||
def _has_mcp_function_name_conflict(self, item_type: str, remote_name: str, local_name: str) -> bool:
|
||||
if self._mcp_reserved_attribute_names is None:
|
||||
self._mcp_reserved_attribute_names = set(dir(self))
|
||||
if local_name not in self._mcp_reserved_attribute_names:
|
||||
return False
|
||||
logger.warning(
|
||||
"Skipping MCP %s '%s' because normalized name '%s' conflicts with an existing plugin attribute.",
|
||||
item_type,
|
||||
remote_name,
|
||||
local_name,
|
||||
)
|
||||
return True
|
||||
|
||||
async def load_prompts(self):
|
||||
"""Load prompts from the MCP server."""
|
||||
try:
|
||||
@@ -472,6 +529,8 @@ class MCPPluginBase:
|
||||
prompt_list = None
|
||||
for prompt in prompt_list.prompts if prompt_list else []:
|
||||
local_name = _normalize_mcp_name(prompt.name)
|
||||
if self._has_mcp_function_name_conflict("prompt", prompt.name, local_name):
|
||||
continue
|
||||
func = kernel_function(name=local_name, description=prompt.description)(
|
||||
partial(self.get_prompt, prompt.name)
|
||||
)
|
||||
@@ -484,9 +543,11 @@ class MCPPluginBase:
|
||||
tool_list = await self.session.list_tools()
|
||||
except Exception:
|
||||
tool_list = None
|
||||
# Create methods with the kernel_function decorator for each tool
|
||||
# Create methods with the kernel_function decorator for each tool
|
||||
for tool in tool_list.tools if tool_list else []:
|
||||
local_name = _normalize_mcp_name(tool.name)
|
||||
if self._has_mcp_function_name_conflict("tool", tool.name, local_name):
|
||||
continue
|
||||
func = kernel_function(name=local_name, description=tool.description)(partial(self.call_tool, tool.name))
|
||||
func.__kernel_function_parameters__ = _get_parameter_dicts_from_mcp_tool(tool)
|
||||
setattr(self, local_name, func)
|
||||
@@ -558,6 +619,7 @@ class MCPStdioPlugin(MCPPluginBase):
|
||||
env: dict[str, str] | None = None,
|
||||
encoding: str | None = None,
|
||||
kernel: Kernel | None = None,
|
||||
sampling_consent_callback: SamplingConsentCallback | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize the MCP stdio plugin.
|
||||
@@ -579,6 +641,10 @@ class MCPStdioPlugin(MCPPluginBase):
|
||||
env: The environment variables to set for the command.
|
||||
encoding: The encoding to use for the command output.
|
||||
kernel: The kernel instance with one or more Chat Completion clients.
|
||||
sampling_consent_callback: Optional callback for approving MCP sampling requests.
|
||||
Receives the plugin name and MCP sampling request params. Return
|
||||
False to deny the request. When omitted, sampling requests are
|
||||
auto-approved and a warning is logged.
|
||||
kwargs: Any extra arguments to pass to the stdio client.
|
||||
|
||||
"""
|
||||
@@ -590,6 +656,7 @@ class MCPStdioPlugin(MCPPluginBase):
|
||||
load_tools=load_tools,
|
||||
load_prompts=load_prompts,
|
||||
request_timeout=request_timeout,
|
||||
sampling_consent_callback=sampling_consent_callback,
|
||||
)
|
||||
self.command = command
|
||||
self.args = args or []
|
||||
@@ -628,6 +695,7 @@ class MCPSsePlugin(MCPPluginBase):
|
||||
timeout: float | None = None,
|
||||
sse_read_timeout: float | None = None,
|
||||
kernel: Kernel | None = None,
|
||||
sampling_consent_callback: SamplingConsentCallback | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize the MCP sse plugin.
|
||||
@@ -650,6 +718,10 @@ class MCPSsePlugin(MCPPluginBase):
|
||||
timeout: The timeout for the request.
|
||||
sse_read_timeout: The timeout for reading from the SSE stream.
|
||||
kernel: The kernel instance with one or more Chat Completion clients.
|
||||
sampling_consent_callback: Optional callback for approving MCP sampling requests.
|
||||
Receives the plugin name and MCP sampling request params. Return
|
||||
False to deny the request. When omitted, sampling requests are
|
||||
auto-approved and a warning is logged.
|
||||
kwargs: Any extra arguments to pass to the sse client.
|
||||
|
||||
"""
|
||||
@@ -661,6 +733,7 @@ class MCPSsePlugin(MCPPluginBase):
|
||||
load_tools=load_tools,
|
||||
load_prompts=load_prompts,
|
||||
request_timeout=request_timeout,
|
||||
sampling_consent_callback=sampling_consent_callback,
|
||||
)
|
||||
self.url = url
|
||||
self.headers = headers or {}
|
||||
@@ -702,6 +775,7 @@ class MCPStreamableHttpPlugin(MCPPluginBase):
|
||||
sse_read_timeout: float | None = None,
|
||||
terminate_on_close: bool | None = None,
|
||||
kernel: Kernel | None = None,
|
||||
sampling_consent_callback: SamplingConsentCallback | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize the MCP streamable http plugin.
|
||||
@@ -725,6 +799,10 @@ class MCPStreamableHttpPlugin(MCPPluginBase):
|
||||
sse_read_timeout: The timeout for reading from the SSE stream.
|
||||
terminate_on_close: Close the transport when the MCP client is terminated.
|
||||
kernel: The kernel instance with one or more Chat Completion clients.
|
||||
sampling_consent_callback: Optional callback for approving MCP sampling requests.
|
||||
Receives the plugin name and MCP sampling request params. Return
|
||||
False to deny the request. When omitted, sampling requests are
|
||||
auto-approved and a warning is logged.
|
||||
kwargs: Any extra arguments to pass to the sse client.
|
||||
"""
|
||||
super().__init__(
|
||||
@@ -735,6 +813,7 @@ class MCPStreamableHttpPlugin(MCPPluginBase):
|
||||
load_tools=load_tools,
|
||||
load_prompts=load_prompts,
|
||||
request_timeout=request_timeout,
|
||||
sampling_consent_callback=sampling_consent_callback,
|
||||
)
|
||||
self.url = url
|
||||
self.headers = headers or {}
|
||||
@@ -775,6 +854,7 @@ class MCPWebsocketPlugin(MCPPluginBase):
|
||||
session: ClientSession | None = None,
|
||||
description: str | None = None,
|
||||
kernel: Kernel | None = None,
|
||||
sampling_consent_callback: SamplingConsentCallback | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize the MCP websocket plugin.
|
||||
@@ -794,6 +874,10 @@ class MCPWebsocketPlugin(MCPPluginBase):
|
||||
session: The session to use for the MCP connection.
|
||||
description: The description of the plugin.
|
||||
kernel: The kernel instance with one or more Chat Completion clients.
|
||||
sampling_consent_callback: Optional callback for approving MCP sampling requests.
|
||||
Receives the plugin name and MCP sampling request params. Return
|
||||
False to deny the request. When omitted, sampling requests are
|
||||
auto-approved and a warning is logged.
|
||||
kwargs: Any extra arguments to pass to the websocket client.
|
||||
|
||||
"""
|
||||
@@ -805,6 +889,7 @@ class MCPWebsocketPlugin(MCPPluginBase):
|
||||
load_tools=load_tools,
|
||||
load_prompts=load_prompts,
|
||||
request_timeout=request_timeout,
|
||||
sampling_consent_callback=sampling_consent_callback,
|
||||
)
|
||||
self.url = url
|
||||
self._client_kwargs = kwargs
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import logging
|
||||
import re
|
||||
from typing import TYPE_CHECKING
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
@@ -89,6 +90,120 @@ async def test_mcp_plugin_session_initialized(plugin_class, plugin_args):
|
||||
assert not mock_session.initialize.called
|
||||
|
||||
|
||||
async def test_mcp_sampling_denied_by_consent_callback():
|
||||
sampling_consent_callback = AsyncMock(return_value=False)
|
||||
plugin = MCPSsePlugin(
|
||||
name="TestMCPPlugin",
|
||||
url="http://localhost:8080/sse",
|
||||
sampling_consent_callback=sampling_consent_callback,
|
||||
)
|
||||
params = types.CreateMessageRequestParams(
|
||||
messages=[types.SamplingMessage(role="user", content=types.TextContent(type="text", text="hello"))],
|
||||
systemPrompt="server instructions",
|
||||
maxTokens=100,
|
||||
)
|
||||
|
||||
result = await plugin.sampling_callback(MagicMock(), params)
|
||||
|
||||
sampling_consent_callback.assert_awaited_once_with("TestMCPPlugin", params)
|
||||
assert isinstance(result, types.ErrorData)
|
||||
assert result.message == "Sampling denied by policy."
|
||||
|
||||
|
||||
async def test_mcp_sampling_consent_callback_error_denies_request(caplog):
|
||||
sampling_consent_callback = AsyncMock(side_effect=RuntimeError("policy failure"))
|
||||
plugin = MCPSsePlugin(
|
||||
name="TestMCPPlugin",
|
||||
url="http://localhost:8080/sse",
|
||||
sampling_consent_callback=sampling_consent_callback,
|
||||
)
|
||||
params = types.CreateMessageRequestParams(
|
||||
messages=[types.SamplingMessage(role="user", content=types.TextContent(type="text", text="hello"))],
|
||||
systemPrompt="server instructions",
|
||||
maxTokens=100,
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.ERROR, logger="semantic_kernel.connectors.mcp"):
|
||||
result = await plugin.sampling_callback(MagicMock(), params)
|
||||
|
||||
sampling_consent_callback.assert_awaited_once_with("TestMCPPlugin", params)
|
||||
assert isinstance(result, types.ErrorData)
|
||||
assert result.message == "Sampling denied by policy."
|
||||
assert "MCP sampling consent callback failed" in caplog.text
|
||||
|
||||
|
||||
async def test_mcp_sampling_without_consent_callback_logs_auto_approve_warning(caplog):
|
||||
plugin = MCPSsePlugin(name="TestMCPPlugin", url="http://localhost:8080/sse")
|
||||
params = types.CreateMessageRequestParams(
|
||||
messages=[types.SamplingMessage(role="user", content=types.TextContent(type="text", text="hello"))],
|
||||
systemPrompt="server instructions",
|
||||
maxTokens=100,
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="semantic_kernel.connectors.mcp"):
|
||||
result = await plugin.sampling_callback(MagicMock(), params)
|
||||
|
||||
assert isinstance(result, types.ErrorData)
|
||||
assert "auto-approved because no sampling consent callback was configured" in caplog.text
|
||||
|
||||
|
||||
async def test_mcp_tool_and_prompt_names_do_not_shadow_plugin_attributes():
|
||||
kernel = MagicMock()
|
||||
plugin = MCPSsePlugin(name="TestMCPPlugin", url="http://localhost:8080/sse", kernel=kernel)
|
||||
session = AsyncMock(spec=ClientSession)
|
||||
session.list_tools.return_value = ListToolsResult(
|
||||
tools=[
|
||||
Tool(name="kernel", description="reserved", inputSchema={}),
|
||||
Tool(name="safe_tool", description="safe", inputSchema={}),
|
||||
]
|
||||
)
|
||||
session.list_prompts.return_value = types.ListPromptsResult(
|
||||
prompts=[
|
||||
types.Prompt(name="session", description="reserved", arguments=[]),
|
||||
types.Prompt(name="safe_prompt", description="safe", arguments=[]),
|
||||
]
|
||||
)
|
||||
plugin.session = session
|
||||
|
||||
await plugin.load_tools()
|
||||
|
||||
assert plugin.kernel is kernel
|
||||
assert hasattr(plugin, "safe_tool")
|
||||
|
||||
await plugin.load_prompts()
|
||||
|
||||
assert plugin.session is session
|
||||
assert hasattr(plugin, "safe_prompt")
|
||||
|
||||
|
||||
async def test_mcp_tool_and_prompt_names_can_reload_existing_mcp_functions():
|
||||
plugin = MCPSsePlugin(name="TestMCPPlugin", url="http://localhost:8080/sse")
|
||||
session = AsyncMock(spec=ClientSession)
|
||||
session.list_tools.side_effect = [
|
||||
ListToolsResult(tools=[Tool(name="safe_tool", description="first tool", inputSchema={})]),
|
||||
ListToolsResult(tools=[Tool(name="safe_tool", description="second tool", inputSchema={})]),
|
||||
]
|
||||
session.list_prompts.side_effect = [
|
||||
types.ListPromptsResult(prompts=[types.Prompt(name="safe_prompt", description="first prompt", arguments=[])]),
|
||||
types.ListPromptsResult(prompts=[types.Prompt(name="safe_prompt", description="second prompt", arguments=[])]),
|
||||
]
|
||||
plugin.session = session
|
||||
|
||||
await plugin.load_tools()
|
||||
first_tool = plugin.safe_tool
|
||||
await plugin.load_tools()
|
||||
|
||||
assert plugin.safe_tool is not first_tool
|
||||
assert plugin.safe_tool.__kernel_function_description__ == "second tool"
|
||||
|
||||
await plugin.load_prompts()
|
||||
first_prompt = plugin.safe_prompt
|
||||
await plugin.load_prompts()
|
||||
|
||||
assert plugin.safe_prompt is not first_prompt
|
||||
assert plugin.safe_prompt.__kernel_function_description__ == "second prompt"
|
||||
|
||||
|
||||
async def test_mcp_plugin_failed_get_session():
|
||||
with (
|
||||
patch("semantic_kernel.connectors.mcp.stdio_client") as mock_stdio_client,
|
||||
|
||||
Reference in New Issue
Block a user