Python: [BREAKING] Add sampling guardrails to MCP tools (#6413)

* Add sampling guardrails to MCP tools

Add approval, token, and request-count controls to the MCP sampling
callback used when an MCPTool is configured with a chat client.

- Add `sampling_approval_callback`, `sampling_max_tokens`, and
  `sampling_max_requests` parameters to `MCPTool` and its
  `MCPStdioTool`, `MCPStreamableHTTPTool`, and `MCPWebsocketTool`
  subclasses, positioned directly after `client`.
- Gate each server-initiated `sampling/createMessage` request behind the
  approval callback, which denies by default when no callback is provided.
- Clamp the requested `maxTokens` to `sampling_max_tokens` and enforce a
  per-session request count via `sampling_max_requests`.
- Log incoming sampling requests at WARNING level (counts only).
- Export `SamplingApprovalCallback` from the public API.
- Add tests, a sample, and documentation updates.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Make sampling denial message context-aware

Distinguish the deny-by-default case (no approval callback configured)
from an explicit denial by a configured `sampling_approval_callback`, so
the returned ErrorData message is accurate for callback-driven denials
and exceptions.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

python/packages/core/tests/core/test_mcp.py

@@ -1813,6 +1813,18 @@ async def test_mcp_tool_message_handler_cancel_and_replace():
     assert len(tool._pending_reload_tasks) == 0
 
 
+def _approve(_params: object) -> bool:
+    """Approving sampling gate used by tests that exercise forwarding behavior."""
+    return True
+
+
+def _make_sampling_response(text: str = "response", model: str = "test-model") -> Mock:
+    mock_response = Mock()
+    mock_response.messages = [Message(role="assistant", contents=[Content.from_text(text)])]
+    mock_response.model = model
+    return mock_response
+
+
 async def test_mcp_tool_sampling_callback_no_client():
     """Test sampling callback error path when no chat client is available."""
     tool = MCPStdioTool(name="test_tool", command="python")
@@ -1828,9 +1840,190 @@ async def test_mcp_tool_sampling_callback_no_client():
     assert "No chat client available" in result.message
 
 
+async def test_mcp_tool_sampling_callback_denies_by_default():
+    """Sampling is denied when no approval callback is configured (safe default)."""
+    tool = MCPStdioTool(name="test_tool", command="python")
+    mock_chat_client = AsyncMock()
+    tool.client = mock_chat_client
+
+    params = Mock()
+    params.messages = []
+    params.maxTokens = 128
+
+    result = await tool.sampling_callback(Mock(), params)
+
+    assert isinstance(result, types.ErrorData)
+    assert result.code == types.INVALID_REQUEST
+    assert "denied" in result.message
+    assert "sampling_approval_callback" in result.message
+    mock_chat_client.get_response.assert_not_called()
+
+
+async def test_mcp_tool_sampling_callback_denied_by_callback():
+    """Sampling is denied when the approval callback returns a falsy value."""
+    tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=lambda params: False)
+    mock_chat_client = AsyncMock()
+    tool.client = mock_chat_client
+
+    params = Mock()
+    params.messages = []
+    params.maxTokens = 128
+
+    result = await tool.sampling_callback(Mock(), params)
+
+    assert isinstance(result, types.ErrorData)
+    assert result.code == types.INVALID_REQUEST
+    assert "denied by the 'sampling_approval_callback'" in result.message
+    mock_chat_client.get_response.assert_not_called()
+
+
+async def test_mcp_tool_sampling_callback_callback_exception_denies():
+    """An approval callback that raises results in denial, not an LLM call."""
+
+    def boom(_params: object) -> bool:
+        raise RuntimeError("approval error")
+
+    tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=boom)
+    mock_chat_client = AsyncMock()
+    tool.client = mock_chat_client
+
+    params = Mock()
+    params.messages = []
+    params.maxTokens = 128
+
+    result = await tool.sampling_callback(Mock(), params)
+
+    assert isinstance(result, types.ErrorData)
+    assert result.code == types.INVALID_REQUEST
+    mock_chat_client.get_response.assert_not_called()
+
+
+async def test_mcp_tool_sampling_callback_async_approval():
+    """An async approval callback that approves allows the request through."""
+
+    async def approve(_params: object) -> bool:
+        return True
+
+    tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=approve)
+    mock_chat_client = AsyncMock()
+    mock_chat_client.get_response.return_value = _make_sampling_response("ok")
+    tool.client = mock_chat_client
+
+    params = Mock()
+    params.messages = [types.PromptMessage(role="user", content=types.TextContent(type="text", text="Hi"))]
+    params.temperature = None
+    params.maxTokens = 100
+    params.stopSequences = None
+    params.systemPrompt = None
+    params.tools = None
+    params.toolChoice = None
+
+    result = await tool.sampling_callback(Mock(), params)
+
+    assert isinstance(result, types.CreateMessageResult)
+    assert result.content.text == "ok"
+    mock_chat_client.get_response.assert_awaited_once()
+
+
+async def test_mcp_tool_sampling_callback_clamps_max_tokens():
+    """An approved request's maxTokens is clamped to sampling_max_tokens."""
+    tool = MCPStdioTool(
+        name="test_tool",
+        command="python",
+        sampling_approval_callback=_approve,
+        sampling_max_tokens=512,
+    )
+    mock_chat_client = AsyncMock()
+    mock_chat_client.get_response.return_value = _make_sampling_response()
+    tool.client = mock_chat_client
+
+    params = Mock()
+    params.messages = [types.PromptMessage(role="user", content=types.TextContent(type="text", text="Hi"))]
+    params.temperature = None
+    params.maxTokens = 1_000_000
+    params.stopSequences = None
+    params.systemPrompt = None
+    params.tools = None
+    params.toolChoice = None
+
+    result = await tool.sampling_callback(Mock(), params)
+
+    assert isinstance(result, types.CreateMessageResult)
+    options = mock_chat_client.get_response.call_args.kwargs.get("options") or {}
+    assert options["max_tokens"] == 512
+
+
+async def test_mcp_tool_sampling_callback_does_not_clamp_under_cap():
+    """A request below the cap keeps its requested maxTokens."""
+    tool = MCPStdioTool(
+        name="test_tool",
+        command="python",
+        sampling_approval_callback=_approve,
+        sampling_max_tokens=512,
+    )
+    mock_chat_client = AsyncMock()
+    mock_chat_client.get_response.return_value = _make_sampling_response()
+    tool.client = mock_chat_client
+
+    params = Mock()
+    params.messages = [types.PromptMessage(role="user", content=types.TextContent(type="text", text="Hi"))]
+    params.temperature = None
+    params.maxTokens = 100
+    params.stopSequences = None
+    params.systemPrompt = None
+    params.tools = None
+    params.toolChoice = None
+
+    result = await tool.sampling_callback(Mock(), params)
+
+    assert isinstance(result, types.CreateMessageResult)
+    options = mock_chat_client.get_response.call_args.kwargs.get("options") or {}
+    assert options["max_tokens"] == 100
+
+
+async def test_mcp_tool_sampling_callback_rate_limited():
+    """Sampling requests beyond sampling_max_requests are rejected per session."""
+    tool = MCPStdioTool(
+        name="test_tool",
+        command="python",
+        sampling_approval_callback=_approve,
+        sampling_max_requests=2,
+    )
+    mock_chat_client = AsyncMock()
+    mock_chat_client.get_response.return_value = _make_sampling_response()
+    tool.client = mock_chat_client
+
+    def make_params() -> Mock:
+        params = Mock()
+        params.messages = [types.PromptMessage(role="user", content=types.TextContent(type="text", text="Hi"))]
+        params.temperature = None
+        params.maxTokens = 100
+        params.stopSequences = None
+        params.systemPrompt = None
+        params.tools = None
+        params.toolChoice = None
+        return params
+
+    first = await tool.sampling_callback(Mock(), make_params())
+    second = await tool.sampling_callback(Mock(), make_params())
+    third = await tool.sampling_callback(Mock(), make_params())
+
+    assert isinstance(first, types.CreateMessageResult)
+    assert isinstance(second, types.CreateMessageResult)
+    assert isinstance(third, types.ErrorData)
+    assert third.code == types.INVALID_REQUEST
+    assert "rate limit" in third.message.lower()
+    assert mock_chat_client.get_response.await_count == 2
+
+    # The counter resets on a session reset.
+    tool._reset_session_state()
+    fourth = await tool.sampling_callback(Mock(), make_params())
+    assert isinstance(fourth, types.CreateMessageResult)
+
+
 async def test_mcp_tool_sampling_callback_chat_client_exception():
     """Test sampling callback when chat client raises exception."""
-    tool = MCPStdioTool(name="test_tool", command="python")
+    tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
 
     # Mock chat client that raises exception
     mock_chat_client = AsyncMock()
@@ -1846,7 +2039,7 @@ async def test_mcp_tool_sampling_callback_chat_client_exception():
     mock_message.content.text = "Test question"
     params.messages = [mock_message]
     params.temperature = None
-    params.maxTokens = None
+    params.maxTokens = 100
     params.stopSequences = None
     params.systemPrompt = None
     params.tools = None
@@ -1863,7 +2056,7 @@ async def test_mcp_tool_sampling_callback_no_valid_content():
     """Test sampling callback when response has no valid content types."""
     from agent_framework import Message
 
-    tool = MCPStdioTool(name="test_tool", command="python")
+    tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
 
     # Mock chat client with response containing only invalid content types
     mock_chat_client = AsyncMock()
@@ -1892,7 +2085,7 @@ async def test_mcp_tool_sampling_callback_no_valid_content():
     mock_message.content.text = "Test question"
     params.messages = [mock_message]
     params.temperature = None
-    params.maxTokens = None
+    params.maxTokens = 100
     params.stopSequences = None
     params.systemPrompt = None
     params.tools = None
@@ -1905,18 +2098,18 @@ async def test_mcp_tool_sampling_callback_no_valid_content():
     assert "Failed to get right content types from the response." in result.message
     mock_chat_client.get_response.assert_awaited_once()
     _, kwargs = mock_chat_client.get_response.await_args
-    assert kwargs["options"] == {"max_tokens": None}
+    assert kwargs["options"] == {"max_tokens": 100}
 
 
 async def test_mcp_tool_sampling_callback_no_response_and_successful_message_creation():
     """Test sampling callback when the chat client returns no response and then valid content."""
-    tool = MCPStdioTool(name="test_tool", command="python")
+    tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
     tool.client = AsyncMock()
 
     params = Mock()
     params.messages = [types.PromptMessage(role="user", content=types.TextContent(type="text", text="Hi"))]
     params.temperature = None
-    params.maxTokens = None
+    params.maxTokens = 100
     params.stopSequences = None
     params.systemPrompt = None
     params.tools = None
@@ -1955,7 +2148,7 @@ async def test_mcp_tool_sampling_callback_forwards_system_prompt():
     """Test sampling callback passes systemPrompt as instructions in options."""
     from agent_framework import Message
 
-    tool = MCPStdioTool(name="test_tool", command="python")
+    tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
 
     mock_chat_client = AsyncMock()
     mock_response = Mock()
@@ -1972,7 +2165,7 @@ async def test_mcp_tool_sampling_callback_forwards_system_prompt():
     mock_message.content.text = "Test question"
     params.messages = [mock_message]
     params.temperature = None
-    params.maxTokens = None
+    params.maxTokens = 100
     params.stopSequences = None
     params.systemPrompt = "You are a helpful assistant"
     params.tools = None
@@ -1990,7 +2183,7 @@ async def test_mcp_tool_sampling_callback_forwards_tools():
     """Test sampling callback converts MCP tools to FunctionTools and passes them in options."""
     from agent_framework import FunctionTool, Message
 
-    tool = MCPStdioTool(name="test_tool", command="python")
+    tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
 
     mock_chat_client = AsyncMock()
     mock_response = Mock()
@@ -2013,7 +2206,7 @@ async def test_mcp_tool_sampling_callback_forwards_tools():
     mock_message.content.text = "Test question"
     params.messages = [mock_message]
     params.temperature = None
-    params.maxTokens = None
+    params.maxTokens = 100
     params.stopSequences = None
     params.systemPrompt = None
     params.tools = [mcp_tool]
@@ -2036,7 +2229,7 @@ async def test_mcp_tool_sampling_callback_forwards_tool_choice():
     """Test sampling callback passes toolChoice mode in options."""
     from agent_framework import Message
 
-    tool = MCPStdioTool(name="test_tool", command="python")
+    tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
 
     mock_chat_client = AsyncMock()
     mock_response = Mock()
@@ -2053,7 +2246,7 @@ async def test_mcp_tool_sampling_callback_forwards_tool_choice():
     mock_message.content.text = "Test question"
     params.messages = [mock_message]
     params.temperature = None
-    params.maxTokens = None
+    params.maxTokens = 100
     params.stopSequences = None
     params.systemPrompt = None
     params.tools = None
@@ -2071,7 +2264,7 @@ async def test_mcp_tool_sampling_callback_forwards_empty_system_prompt():
     """Test sampling callback forwards empty string systemPrompt as instructions."""
     from agent_framework import Message
 
-    tool = MCPStdioTool(name="test_tool", command="python")
+    tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
 
     mock_chat_client = AsyncMock()
     mock_response = Mock()
@@ -2088,7 +2281,7 @@ async def test_mcp_tool_sampling_callback_forwards_empty_system_prompt():
     mock_message.content.text = "Test question"
     params.messages = [mock_message]
     params.temperature = None
-    params.maxTokens = None
+    params.maxTokens = 100
     params.stopSequences = None
     params.systemPrompt = ""
     params.tools = None
@@ -2106,7 +2299,7 @@ async def test_mcp_tool_sampling_callback_forwards_empty_tools_list():
     """Test sampling callback forwards empty tools list in options."""
     from agent_framework import Message
 
-    tool = MCPStdioTool(name="test_tool", command="python")
+    tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
 
     mock_chat_client = AsyncMock()
     mock_response = Mock()
@@ -2123,7 +2316,7 @@ async def test_mcp_tool_sampling_callback_forwards_empty_tools_list():
     mock_message.content.text = "Test question"
     params.messages = [mock_message]
     params.temperature = None
-    params.maxTokens = None
+    params.maxTokens = 100
     params.stopSequences = None
     params.systemPrompt = None
     params.tools = []
@@ -2141,7 +2334,7 @@ async def test_mcp_tool_sampling_callback_forwards_generation_params_in_options(
     """Test sampling callback passes temperature, max_tokens, and stop in options."""
     from agent_framework import Message
 
-    tool = MCPStdioTool(name="test_tool", command="python")
+    tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
 
     mock_chat_client = AsyncMock()
     mock_response = Mock()
@@ -2182,7 +2375,7 @@ async def test_mcp_tool_sampling_callback_omits_temperature_when_none():
     """Test sampling callback does not set temperature in options when it is None."""
     from agent_framework import Message
 
-    tool = MCPStdioTool(name="test_tool", command="python")
+    tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
 
     mock_chat_client = AsyncMock()
     mock_response = Mock()
@@ -2219,7 +2412,7 @@ async def test_mcp_tool_sampling_callback_always_passes_max_tokens():
     """Test sampling callback always sets max_tokens in options since maxTokens is a required int field."""
     from agent_framework import Message
 
-    tool = MCPStdioTool(name="test_tool", command="python")
+    tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
 
     mock_chat_client = AsyncMock()
     mock_response = Mock()