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:
Evan Mattson
2026-05-14 00:46:19 +00:00
committed by GitHub
co-authored by Copilot
parent 32abf99e29
commit a6c68af594
2 changed files with 205 additions and 5 deletions
+90 -5
View File
@@ -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,