From bbeec89f7ea7ba0746156a4b6ccae09b6322e46a Mon Sep 17 00:00:00 2001 From: Copilot Date: Wed, 18 Mar 2026 13:22:44 +0000 Subject: [PATCH 1/9] Support MCP sampling tools capability (#4625) Forward systemPrompt, tools, and toolChoice from MCP sampling requests to the chat client's get_response() call. Also advertise the sampling.tools capability to MCP servers when a client is configured. - Pass SamplingCapability with tools support to ClientSession - Convert systemPrompt to instructions in options - Convert MCP Tool objects to FunctionTool instances for options - Map MCP ToolChoice.mode to tool_choice in options - Add tests for all new behaviors and update existing sampling tests Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/core/agent_framework/_mcp.py | 23 +++ python/packages/core/tests/core/test_mcp.py | 184 +++++++++++++++++++ 2 files changed, 207 insertions(+) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 062df5491c..0ba4e207f1 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -649,6 +649,11 @@ async def _connect_on_owner(self, *, reset: bool = False) -> None: error_msg = f"Failed to connect to MCP server: {ex}" raise ToolException(error_msg, inner_exception=ex) from ex try: + sampling_capabilities = None + if self.client is not None: + sampling_capabilities = types.SamplingCapability( + tools=types.SamplingToolsCapability(), + ) session = await self._exit_stack.enter_async_context( ClientSession( read_stream=transport[0], @@ -659,6 +664,7 @@ async def _connect_on_owner(self, *, reset: bool = False) -> None: message_handler=self.message_handler, logging_callback=self.logging_callback, sampling_callback=self.sampling_callback, + sampling_capabilities=sampling_capabilities, ) ) except Exception as ex: @@ -732,12 +738,29 @@ async def sampling_callback( messages: list[Message] = [] for msg in params.messages: messages.append(_parse_message_from_mcp(msg)) + + options: dict[str, Any] = {} + if params.systemPrompt: + options["instructions"] = params.systemPrompt + if params.tools: + options["tools"] = [ + FunctionTool( + name=tool.name, + description=tool.description or "", + input_model=tool.inputSchema, + ) + for tool in params.tools + ] + if params.toolChoice is not None and params.toolChoice.mode is not None: + options["tool_choice"] = params.toolChoice.mode + try: response = await self.client.get_response( messages, temperature=params.temperature, max_tokens=params.maxTokens, stop=params.stopSequences, + options=options or None, ) except Exception as ex: return types.ErrorData( diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index fa9e1130f0..dc6716705c 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -1562,6 +1562,9 @@ async def test_mcp_tool_sampling_callback_chat_client_exception(): params.temperature = None params.maxTokens = None params.stopSequences = None + params.systemPrompt = None + params.tools = None + params.toolChoice = None result = await tool.sampling_callback(Mock(), params) @@ -1605,6 +1608,9 @@ async def test_mcp_tool_sampling_callback_no_valid_content(): params.temperature = None params.maxTokens = None params.stopSequences = None + params.systemPrompt = None + params.tools = None + params.toolChoice = None result = await tool.sampling_callback(Mock(), params) @@ -1613,6 +1619,184 @@ async def test_mcp_tool_sampling_callback_no_valid_content(): assert "Failed to get right content types from the response." in result.message +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") + + mock_chat_client = AsyncMock() + mock_response = Mock() + mock_response.messages = [ + Message(role="assistant", contents=[Content.from_text("response")]) + ] + mock_response.model_id = "test-model" + mock_chat_client.get_response.return_value = mock_response + + tool.client = mock_chat_client + + params = Mock() + mock_message = Mock() + mock_message.role = "user" + mock_message.content = Mock() + mock_message.content.text = "Test question" + params.messages = [mock_message] + params.temperature = None + params.maxTokens = None + params.stopSequences = None + params.systemPrompt = "You are a helpful assistant" + params.tools = None + params.toolChoice = None + + result = await tool.sampling_callback(Mock(), params) + + assert isinstance(result, types.CreateMessageResult) + call_kwargs = mock_chat_client.get_response.call_args + options = call_kwargs.kwargs.get("options") or {} + assert options.get("instructions") == "You are a helpful assistant" + + +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") + + mock_chat_client = AsyncMock() + mock_response = Mock() + mock_response.messages = [ + Message(role="assistant", contents=[Content.from_text("response")]) + ] + mock_response.model_id = "test-model" + mock_chat_client.get_response.return_value = mock_response + + tool.client = mock_chat_client + + mcp_tool = types.Tool( + name="get_weather", + description="Get weather", + inputSchema={"type": "object", "properties": {"city": {"type": "string"}}}, + ) + + params = Mock() + mock_message = Mock() + mock_message.role = "user" + mock_message.content = Mock() + mock_message.content.text = "Test question" + params.messages = [mock_message] + params.temperature = None + params.maxTokens = None + params.stopSequences = None + params.systemPrompt = None + params.tools = [mcp_tool] + params.toolChoice = None + + result = await tool.sampling_callback(Mock(), params) + + assert isinstance(result, types.CreateMessageResult) + call_kwargs = mock_chat_client.get_response.call_args + options = call_kwargs.kwargs.get("options") or {} + tools = options.get("tools") + assert tools is not None + assert len(tools) == 1 + assert isinstance(tools[0], FunctionTool) + assert tools[0].name == "get_weather" + assert tools[0].description == "Get weather" + + +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") + + mock_chat_client = AsyncMock() + mock_response = Mock() + mock_response.messages = [ + Message(role="assistant", contents=[Content.from_text("response")]) + ] + mock_response.model_id = "test-model" + mock_chat_client.get_response.return_value = mock_response + + tool.client = mock_chat_client + + params = Mock() + mock_message = Mock() + mock_message.role = "user" + mock_message.content = Mock() + mock_message.content.text = "Test question" + params.messages = [mock_message] + params.temperature = None + params.maxTokens = None + params.stopSequences = None + params.systemPrompt = None + params.tools = None + params.toolChoice = types.ToolChoice(mode="required") + + result = await tool.sampling_callback(Mock(), params) + + assert isinstance(result, types.CreateMessageResult) + call_kwargs = mock_chat_client.get_response.call_args + options = call_kwargs.kwargs.get("options") or {} + assert options.get("tool_choice") == "required" + + +async def test_connect_creates_session_with_sampling_capabilities(): + """Test connect() passes sampling_capabilities to ClientSession when client is set.""" + tool = MCPStdioTool(name="test", command="test-command", load_tools=False, load_prompts=False) + tool.client = Mock() + + mock_transport = (Mock(), Mock()) + mock_context_manager = Mock() + mock_context_manager.__aenter__ = AsyncMock(return_value=mock_transport) + mock_context_manager.__aexit__ = AsyncMock(return_value=None) + tool.get_mcp_client = Mock(return_value=mock_context_manager) + + with patch("agent_framework._mcp.ClientSession") as mock_session_class: + mock_session = AsyncMock() + mock_session._request_id = 1 + + session_cm = AsyncMock() + session_cm.__aenter__ = AsyncMock(return_value=mock_session) + session_cm.__aexit__ = AsyncMock(return_value=None) + mock_session_class.return_value = session_cm + + await tool.connect() + + call_kwargs = mock_session_class.call_args.kwargs + sampling_caps = call_kwargs.get("sampling_capabilities") + assert sampling_caps is not None + assert isinstance(sampling_caps, types.SamplingCapability) + assert sampling_caps.tools is not None + assert isinstance(sampling_caps.tools, types.SamplingToolsCapability) + + +async def test_connect_no_sampling_capabilities_without_client(): + """Test connect() does not pass sampling_capabilities when no client is set.""" + tool = MCPStdioTool(name="test", command="test-command", load_tools=False, load_prompts=False) + # No client set + + mock_transport = (Mock(), Mock()) + mock_context_manager = Mock() + mock_context_manager.__aenter__ = AsyncMock(return_value=mock_transport) + mock_context_manager.__aexit__ = AsyncMock(return_value=None) + tool.get_mcp_client = Mock(return_value=mock_context_manager) + + with patch("agent_framework._mcp.ClientSession") as mock_session_class: + mock_session = AsyncMock() + mock_session._request_id = 1 + + session_cm = AsyncMock() + session_cm.__aenter__ = AsyncMock(return_value=mock_session) + session_cm.__aexit__ = AsyncMock(return_value=None) + mock_session_class.return_value = session_cm + + await tool.connect() + + call_kwargs = mock_session_class.call_args.kwargs + assert call_kwargs.get("sampling_capabilities") is None + + # Test error handling in connect() method From 2ad3fbd394b620a2d1c5f7c670753eaa657d147f Mon Sep 17 00:00:00 2001 From: Copilot Date: Wed, 18 Mar 2026 13:23:09 +0000 Subject: [PATCH 2/9] Apply pre-commit auto-fixes --- python/packages/core/tests/core/test_mcp.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index dc6716705c..a384d2fcf5 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -1627,9 +1627,7 @@ async def test_mcp_tool_sampling_callback_forwards_system_prompt(): mock_chat_client = AsyncMock() mock_response = Mock() - mock_response.messages = [ - Message(role="assistant", contents=[Content.from_text("response")]) - ] + mock_response.messages = [Message(role="assistant", contents=[Content.from_text("response")])] mock_response.model_id = "test-model" mock_chat_client.get_response.return_value = mock_response @@ -1664,9 +1662,7 @@ async def test_mcp_tool_sampling_callback_forwards_tools(): mock_chat_client = AsyncMock() mock_response = Mock() - mock_response.messages = [ - Message(role="assistant", contents=[Content.from_text("response")]) - ] + mock_response.messages = [Message(role="assistant", contents=[Content.from_text("response")])] mock_response.model_id = "test-model" mock_chat_client.get_response.return_value = mock_response @@ -1712,9 +1708,7 @@ async def test_mcp_tool_sampling_callback_forwards_tool_choice(): mock_chat_client = AsyncMock() mock_response = Mock() - mock_response.messages = [ - Message(role="assistant", contents=[Content.from_text("response")]) - ] + mock_response.messages = [Message(role="assistant", contents=[Content.from_text("response")])] mock_response.model_id = "test-model" mock_chat_client.get_response.return_value = mock_response From fa3d3c4e5c844f40f5daadb3b5ad2324f3e45c09 Mon Sep 17 00:00:00 2001 From: Copilot Date: Wed, 18 Mar 2026 13:29:33 +0000 Subject: [PATCH 3/9] Fix #4625: Support MCP sampling tool with proper typing and structured content - Fix mypy error by typing sampling callback options as ChatOptions[None] instead of dict[str, Any], and importing ChatOptions from _types - Handle structuredContent from CallToolResult in _parse_tool_result_from_mcp, serializing it as JSON text Content when present - Add tests for structuredContent parsing (with and without regular content) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/core/agent_framework/_mcp.py | 6 +++- python/packages/core/tests/core/test_mcp.py | 33 ++++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 0ba4e207f1..a077d7af36 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -29,6 +29,7 @@ from ._tools import FunctionTool from ._types import ( + ChatOptions, Content, Message, ) @@ -192,6 +193,9 @@ def _parse_tool_result_from_mcp( case _: result.append(Content.from_text(str(item))) + if mcp_type.structuredContent is not None: + result.append(Content.from_text(json.dumps(mcp_type.structuredContent))) + if not result: result.append(Content.from_text("null")) return result @@ -739,7 +743,7 @@ async def sampling_callback( for msg in params.messages: messages.append(_parse_message_from_mcp(msg)) - options: dict[str, Any] = {} + options: ChatOptions[None] = {} if params.systemPrompt: options["instructions"] = params.systemPrompt if params.tools: diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index a384d2fcf5..ec4f050fe0 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. # type: ignore[reportPrivateUsage] +import json import logging import os from contextlib import _AsyncGeneratorContextManager # type: ignore @@ -246,6 +247,38 @@ def test_parse_tool_result_from_mcp_blob_plain_base64(): assert "dGVzdCBkYXRh" in result[0].uri +def test_parse_tool_result_from_mcp_structured_content(): + """Test that structuredContent from CallToolResult is included as JSON text.""" + structured = {"name": "Pasta Carbonara", "ingredients": ["pasta", "eggs", "cheese"]} + mcp_result = types.CallToolResult( + content=[types.TextContent(type="text", text="Here is a recipe")], + structuredContent=structured, + ) + result = _parse_tool_result_from_mcp(mcp_result) + + assert isinstance(result, list) + assert len(result) == 2 + assert result[0].type == "text" + assert result[0].text == "Here is a recipe" + assert result[1].type == "text" + assert result[1].text == json.dumps(structured) + + +def test_parse_tool_result_from_mcp_structured_content_only(): + """Test that structuredContent alone (no regular content) produces a text Content.""" + structured = {"temperature": 72, "unit": "F"} + mcp_result = types.CallToolResult( + content=[], + structuredContent=structured, + ) + result = _parse_tool_result_from_mcp(mcp_result) + + assert isinstance(result, list) + assert len(result) == 1 + assert result[0].type == "text" + assert result[0].text == json.dumps(structured) + + def test_mcp_content_types_to_ai_content_text(): """Test conversion of MCP text content to AI content.""" mcp_content = types.TextContent(type="text", text="Sample text") From 598b5e4d436dc82a03211c320fd66d48bbc3aa95 Mon Sep 17 00:00:00 2001 From: Copilot Date: Wed, 18 Mar 2026 13:44:28 +0000 Subject: [PATCH 4/9] Fix lint: add author to TODO comment Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/core/agent_framework/_mcp.py | 15 ++++++---- python/packages/core/tests/core/test_mcp.py | 29 ++++++++++++++++---- 2 files changed, 34 insertions(+), 10 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index a077d7af36..cbdae0f021 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -147,7 +147,10 @@ def _parse_tool_result_from_mcp( Converts each content item in the MCP result to its appropriate Content form. Text items become ``Content(type="text")`` and media - items (images, audio) are preserved as rich Content. + items (images, audio) are preserved as rich Content. When no content + items are produced but ``structuredContent`` is present on the result, + the structured payload is serialised as JSON text so it is still + surfaced to the caller. Args: mcp_type: The MCP CallToolResult object to convert. @@ -193,8 +196,8 @@ def _parse_tool_result_from_mcp( case _: result.append(Content.from_text(str(item))) - if mcp_type.structuredContent is not None: - result.append(Content.from_text(json.dumps(mcp_type.structuredContent))) + if not result and mcp_type.structuredContent is not None: + result.append(Content.from_text(json.dumps(mcp_type.structuredContent, default=str))) if not result: result.append(Content.from_text("null")) @@ -743,10 +746,12 @@ async def sampling_callback( for msg in params.messages: messages.append(_parse_message_from_mcp(msg)) + # TODO(Copilot): forward server's expected result schema as response_format + # when result_type support is added to ChatOptions. options: ChatOptions[None] = {} - if params.systemPrompt: + if params.systemPrompt is not None: options["instructions"] = params.systemPrompt - if params.tools: + if params.tools is not None: options["tools"] = [ FunctionTool( name=tool.name, diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index ec4f050fe0..d3706c13fe 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -248,7 +248,7 @@ def test_parse_tool_result_from_mcp_blob_plain_base64(): def test_parse_tool_result_from_mcp_structured_content(): - """Test that structuredContent from CallToolResult is included as JSON text.""" + """Test that structuredContent is ignored when content items are present.""" structured = {"name": "Pasta Carbonara", "ingredients": ["pasta", "eggs", "cheese"]} mcp_result = types.CallToolResult( content=[types.TextContent(type="text", text="Here is a recipe")], @@ -257,11 +257,9 @@ def test_parse_tool_result_from_mcp_structured_content(): result = _parse_tool_result_from_mcp(mcp_result) assert isinstance(result, list) - assert len(result) == 2 + assert len(result) == 1 assert result[0].type == "text" assert result[0].text == "Here is a recipe" - assert result[1].type == "text" - assert result[1].text == json.dumps(structured) def test_parse_tool_result_from_mcp_structured_content_only(): @@ -276,7 +274,28 @@ def test_parse_tool_result_from_mcp_structured_content_only(): assert isinstance(result, list) assert len(result) == 1 assert result[0].type == "text" - assert result[0].text == json.dumps(structured) + assert json.loads(result[0].text) == structured + + +def test_parse_tool_result_from_mcp_structured_content_nested(): + """Test that structuredContent with nested complex types serialises correctly.""" + structured = { + "recipe": { + "name": "Pasta Carbonara", + "ingredients": [{"item": "pasta", "amount": 200}, {"item": "eggs", "amount": 3}], + "metadata": {"origin": "Italy", "tags": ["quick", "classic"]}, + } + } + mcp_result = types.CallToolResult( + content=[], + structuredContent=structured, + ) + result = _parse_tool_result_from_mcp(mcp_result) + + assert isinstance(result, list) + assert len(result) == 1 + assert result[0].type == "text" + assert json.loads(result[0].text) == structured def test_mcp_content_types_to_ai_content_text(): From 4322e90b4a682fab5e0eab51d30b5133eb264658 Mon Sep 17 00:00:00 2001 From: Copilot Date: Wed, 18 Mar 2026 13:52:44 +0000 Subject: [PATCH 5/9] Address review feedback for #4625: remove default=str, add edge-case tests - Remove default=str from json.dumps for structuredContent to fail fast on non-JSON-serializable values instead of silently converting - Add test for non-JSON-serializable structuredContent (TypeError) - Add tests for empty systemPrompt ('') and empty tools list ([]) edge cases in sampling callback - Expand TODO comment noting list[Content] return type constraint for future result_type support Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/core/agent_framework/_mcp.py | 5 +- python/packages/core/tests/core/test_mcp.py | 83 ++++++++++++++++++++ 2 files changed, 87 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index cbdae0f021..edb6eb4ee8 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -197,7 +197,7 @@ def _parse_tool_result_from_mcp( result.append(Content.from_text(str(item))) if not result and mcp_type.structuredContent is not None: - result.append(Content.from_text(json.dumps(mcp_type.structuredContent, default=str))) + result.append(Content.from_text(json.dumps(mcp_type.structuredContent))) if not result: result.append(Content.from_text("null")) @@ -748,6 +748,9 @@ async def sampling_callback( # TODO(Copilot): forward server's expected result schema as response_format # when result_type support is added to ChatOptions. + # Note: _parse_tool_result_from_mcp returns list[Content] with no structured-data + # channel, so retrofitting typed results will require a parallel code path or + # richer return type. options: ChatOptions[None] = {} if params.systemPrompt is not None: options["instructions"] = params.systemPrompt diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index d3706c13fe..90b8d9da16 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -298,6 +298,19 @@ def test_parse_tool_result_from_mcp_structured_content_nested(): assert json.loads(result[0].text) == structured +def test_parse_tool_result_from_mcp_structured_content_non_serializable(): + """Test that structuredContent with non-JSON-serializable values raises TypeError.""" + from datetime import datetime + + structured = {"timestamp": datetime(2025, 1, 1)} + mcp_result = types.CallToolResult( + content=[], + structuredContent=structured, + ) + with pytest.raises(TypeError): + _parse_tool_result_from_mcp(mcp_result) + + def test_mcp_content_types_to_ai_content_text(): """Test conversion of MCP text content to AI content.""" mcp_content = types.TextContent(type="text", text="Sample text") @@ -1787,6 +1800,76 @@ async def test_mcp_tool_sampling_callback_forwards_tool_choice(): assert options.get("tool_choice") == "required" +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") + + mock_chat_client = AsyncMock() + mock_response = Mock() + mock_response.messages = [Message(role="assistant", contents=[Content.from_text("response")])] + mock_response.model_id = "test-model" + mock_chat_client.get_response.return_value = mock_response + + tool.client = mock_chat_client + + params = Mock() + mock_message = Mock() + mock_message.role = "user" + mock_message.content = Mock() + mock_message.content.text = "Test question" + params.messages = [mock_message] + params.temperature = None + params.maxTokens = None + params.stopSequences = None + params.systemPrompt = "" + params.tools = None + params.toolChoice = None + + result = await tool.sampling_callback(Mock(), params) + + assert isinstance(result, types.CreateMessageResult) + call_kwargs = mock_chat_client.get_response.call_args + options = call_kwargs.kwargs.get("options") or {} + assert options.get("instructions") == "" + + +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") + + mock_chat_client = AsyncMock() + mock_response = Mock() + mock_response.messages = [Message(role="assistant", contents=[Content.from_text("response")])] + mock_response.model_id = "test-model" + mock_chat_client.get_response.return_value = mock_response + + tool.client = mock_chat_client + + params = Mock() + mock_message = Mock() + mock_message.role = "user" + mock_message.content = Mock() + mock_message.content.text = "Test question" + params.messages = [mock_message] + params.temperature = None + params.maxTokens = None + params.stopSequences = None + params.systemPrompt = None + params.tools = [] + params.toolChoice = None + + result = await tool.sampling_callback(Mock(), params) + + assert isinstance(result, types.CreateMessageResult) + call_kwargs = mock_chat_client.get_response.call_args + options = call_kwargs.kwargs.get("options") or {} + assert options.get("tools") == [] + + async def test_connect_creates_session_with_sampling_capabilities(): """Test connect() passes sampling_capabilities to ClientSession when client is set.""" tool = MCPStdioTool(name="test", command="test-command", load_tools=False, load_prompts=False) From 4551a1e25625c351a9dd80fb86c70b8eb54dfe1f Mon Sep 17 00:00:00 2001 From: Copilot Date: Wed, 18 Mar 2026 13:53:31 +0000 Subject: [PATCH 6/9] Sanitize sampling callback error to avoid leaking internals (#4625) Log exception details at DEBUG level instead of including them in the ErrorData message returned to the MCP server, which may be untrusted. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/core/agent_framework/_mcp.py | 3 ++- python/packages/core/tests/core/test_mcp.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index edb6eb4ee8..92e035f2bc 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -775,9 +775,10 @@ async def sampling_callback( options=options or None, ) except Exception as ex: + logger.debug("Sampling callback error: %s", ex, exc_info=True) return types.ErrorData( code=types.INTERNAL_ERROR, - message=f"Failed to get chat message content: {ex}", + message="Failed to get chat message content.", ) if not response or not response.messages: return types.ErrorData( diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 90b8d9da16..31214785b6 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -1635,7 +1635,7 @@ async def test_mcp_tool_sampling_callback_chat_client_exception(): assert isinstance(result, types.ErrorData) assert result.code == types.INTERNAL_ERROR - assert "Failed to get chat message content: Chat client error" in result.message + assert "Failed to get chat message content" in result.message async def test_mcp_tool_sampling_callback_no_valid_content(): From 8a89de90dea218b9282abf76f8a835cbaa1f8d73 Mon Sep 17 00:00:00 2001 From: Copilot Date: Wed, 18 Mar 2026 14:36:41 +0000 Subject: [PATCH 7/9] Address review feedback for #4625: move params to options, restore error info - Remove stale TODO comment about response_format (ChatOptions already has it) - Restore {ex} in sampling callback error message for useful debugging info - Set structuredContent as additional_property on Content for structured access - Move temperature, max_tokens, stop into options dict (not top-level kwargs) - Only set temperature when provided (not all models support it) - Add tests for generation params in options and temperature omission Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/core/agent_framework/_mcp.py | 23 +++--- python/packages/core/tests/core/test_mcp.py | 79 +++++++++++++++++++- 2 files changed, 91 insertions(+), 11 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 92e035f2bc..49127f9658 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -197,7 +197,12 @@ def _parse_tool_result_from_mcp( result.append(Content.from_text(str(item))) if not result and mcp_type.structuredContent is not None: - result.append(Content.from_text(json.dumps(mcp_type.structuredContent))) + result.append( + Content.from_text( + json.dumps(mcp_type.structuredContent), + additional_properties={"structured_content": mcp_type.structuredContent}, + ) + ) if not result: result.append(Content.from_text("null")) @@ -746,11 +751,6 @@ async def sampling_callback( for msg in params.messages: messages.append(_parse_message_from_mcp(msg)) - # TODO(Copilot): forward server's expected result schema as response_format - # when result_type support is added to ChatOptions. - # Note: _parse_tool_result_from_mcp returns list[Content] with no structured-data - # channel, so retrofitting typed results will require a parallel code path or - # richer return type. options: ChatOptions[None] = {} if params.systemPrompt is not None: options["instructions"] = params.systemPrompt @@ -766,19 +766,22 @@ async def sampling_callback( if params.toolChoice is not None and params.toolChoice.mode is not None: options["tool_choice"] = params.toolChoice.mode + if params.temperature is not None: + options["temperature"] = params.temperature + options["max_tokens"] = params.maxTokens + if params.stopSequences is not None: + options["stop"] = params.stopSequences + try: response = await self.client.get_response( messages, - temperature=params.temperature, - max_tokens=params.maxTokens, - stop=params.stopSequences, options=options or None, ) except Exception as ex: logger.debug("Sampling callback error: %s", ex, exc_info=True) return types.ErrorData( code=types.INTERNAL_ERROR, - message="Failed to get chat message content.", + message=f"Failed to get chat message content: {ex}", ) if not response or not response.messages: return types.ErrorData( diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 31214785b6..f3453858b6 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -275,6 +275,8 @@ def test_parse_tool_result_from_mcp_structured_content_only(): assert len(result) == 1 assert result[0].type == "text" assert json.loads(result[0].text) == structured + assert result[0].additional_properties is not None + assert result[0].additional_properties["structured_content"] == structured def test_parse_tool_result_from_mcp_structured_content_nested(): @@ -1870,7 +1872,82 @@ async def test_mcp_tool_sampling_callback_forwards_empty_tools_list(): assert options.get("tools") == [] -async def test_connect_creates_session_with_sampling_capabilities(): +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") + + mock_chat_client = AsyncMock() + mock_response = Mock() + mock_response.messages = [Message(role="assistant", contents=[Content.from_text("response")])] + mock_response.model_id = "test-model" + mock_chat_client.get_response.return_value = mock_response + + tool.client = mock_chat_client + + params = Mock() + mock_message = Mock() + mock_message.role = "user" + mock_message.content = Mock() + mock_message.content.text = "Test question" + params.messages = [mock_message] + params.temperature = 0.7 + params.maxTokens = 256 + params.stopSequences = ["STOP"] + params.systemPrompt = None + params.tools = None + params.toolChoice = None + + result = await tool.sampling_callback(Mock(), params) + + assert isinstance(result, types.CreateMessageResult) + call_kwargs = mock_chat_client.get_response.call_args + options = call_kwargs.kwargs.get("options") or {} + assert options.get("temperature") == 0.7 + assert options.get("max_tokens") == 256 + assert options.get("stop") == ["STOP"] + # These should not be passed as top-level kwargs + assert "temperature" not in call_kwargs.kwargs + assert "max_tokens" not in call_kwargs.kwargs + assert "stop" not in call_kwargs.kwargs + + +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") + + mock_chat_client = AsyncMock() + mock_response = Mock() + mock_response.messages = [Message(role="assistant", contents=[Content.from_text("response")])] + mock_response.model_id = "test-model" + mock_chat_client.get_response.return_value = mock_response + + tool.client = mock_chat_client + + params = Mock() + mock_message = Mock() + mock_message.role = "user" + mock_message.content = Mock() + mock_message.content.text = "Test question" + params.messages = [mock_message] + 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) + call_kwargs = mock_chat_client.get_response.call_args + options = call_kwargs.kwargs.get("options") or {} + assert "temperature" not in options + assert options.get("max_tokens") == 100 + assert "stop" not in options """Test connect() passes sampling_capabilities to ClientSession when client is set.""" tool = MCPStdioTool(name="test", command="test-command", load_tools=False, load_prompts=False) tool.client = Mock() From f13bf56cfa761e9a05fa309666e0cbefb71dcbb6 Mon Sep 17 00:00:00 2001 From: Copilot Date: Thu, 19 Mar 2026 10:59:56 +0000 Subject: [PATCH 8/9] Fix MCP sampling callback and structured content error handling (#4625) - Guard max_tokens like temperature: only set when not None, so options can properly evaluate to None when all params are absent - Wrap json.dumps of structuredContent in try/except to fall back to str() for non-serializable values instead of propagating TypeError - Extract test_connect_sampling_capabilities_with_client into its own test function so pytest can discover it independently - Add test for max_tokens=None omission from options - Update structured content non-serializable test to expect fallback Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/core/agent_framework/_mcp.py | 9 +++- python/packages/core/tests/core/test_mcp.py | 46 ++++++++++++++++++-- 2 files changed, 50 insertions(+), 5 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 49127f9658..7ba1ac0268 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -197,9 +197,13 @@ def _parse_tool_result_from_mcp( result.append(Content.from_text(str(item))) if not result and mcp_type.structuredContent is not None: + try: + text = json.dumps(mcp_type.structuredContent) + except (TypeError, ValueError): + text = str(mcp_type.structuredContent) result.append( Content.from_text( - json.dumps(mcp_type.structuredContent), + text, additional_properties={"structured_content": mcp_type.structuredContent}, ) ) @@ -768,7 +772,8 @@ async def sampling_callback( if params.temperature is not None: options["temperature"] = params.temperature - options["max_tokens"] = params.maxTokens + if params.maxTokens is not None: + options["max_tokens"] = params.maxTokens if params.stopSequences is not None: options["stop"] = params.stopSequences diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index f3453858b6..14e91d6786 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -301,7 +301,7 @@ def test_parse_tool_result_from_mcp_structured_content_nested(): def test_parse_tool_result_from_mcp_structured_content_non_serializable(): - """Test that structuredContent with non-JSON-serializable values raises TypeError.""" + """Test that structuredContent with non-JSON-serializable values falls back to str().""" from datetime import datetime structured = {"timestamp": datetime(2025, 1, 1)} @@ -309,8 +309,10 @@ def test_parse_tool_result_from_mcp_structured_content_non_serializable(): content=[], structuredContent=structured, ) - with pytest.raises(TypeError): - _parse_tool_result_from_mcp(mcp_result) + result = _parse_tool_result_from_mcp(mcp_result) + assert len(result) == 1 + assert result[0].text == str(structured) + assert result[0].additional_properties["structured_content"] == structured def test_mcp_content_types_to_ai_content_text(): @@ -1948,6 +1950,44 @@ async def test_mcp_tool_sampling_callback_omits_temperature_when_none(): assert "temperature" not in options assert options.get("max_tokens") == 100 assert "stop" not in options + + +async def test_mcp_tool_sampling_callback_omits_max_tokens_when_none(): + """Test sampling callback does not set max_tokens in options when it is None.""" + from agent_framework import Message + + tool = MCPStdioTool(name="test_tool", command="python") + + mock_chat_client = AsyncMock() + mock_response = Mock() + mock_response.messages = [Message(role="assistant", contents=[Content.from_text("response")])] + mock_response.model_id = "test-model" + mock_chat_client.get_response.return_value = mock_response + + tool.client = mock_chat_client + + params = Mock() + mock_message = Mock() + mock_message.role = "user" + mock_message.content = Mock() + mock_message.content.text = "Test question" + params.messages = [mock_message] + params.temperature = None + params.maxTokens = None + params.stopSequences = None + params.systemPrompt = None + params.tools = None + params.toolChoice = None + + result = await tool.sampling_callback(Mock(), params) + + assert isinstance(result, types.CreateMessageResult) + call_kwargs = mock_chat_client.get_response.call_args + options = call_kwargs.kwargs.get("options") + assert options is None or "max_tokens" not in options + + +async def test_connect_sampling_capabilities_with_client(): """Test connect() passes sampling_capabilities to ClientSession when client is set.""" tool = MCPStdioTool(name="test", command="test-command", load_tools=False, load_prompts=False) tool.client = Mock() From e08036a2acd669841ed95cc9a74fea6a29a97a31 Mon Sep 17 00:00:00 2001 From: Copilot Date: Thu, 19 Mar 2026 11:03:47 +0000 Subject: [PATCH 9/9] Address review feedback for #4625: review comment fixes --- python/packages/core/agent_framework/_mcp.py | 3 +-- python/packages/core/tests/core/test_mcp.py | 10 +++++----- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 7ba1ac0268..875e0fa5ca 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -772,8 +772,7 @@ async def sampling_callback( if params.temperature is not None: options["temperature"] = params.temperature - if params.maxTokens is not None: - options["max_tokens"] = params.maxTokens + options["max_tokens"] = params.maxTokens if params.stopSequences is not None: options["stop"] = params.stopSequences diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 14e91d6786..be9ec718cc 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -1952,8 +1952,8 @@ async def test_mcp_tool_sampling_callback_omits_temperature_when_none(): assert "stop" not in options -async def test_mcp_tool_sampling_callback_omits_max_tokens_when_none(): - """Test sampling callback does not set max_tokens in options when it is None.""" +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") @@ -1973,7 +1973,7 @@ async def test_mcp_tool_sampling_callback_omits_max_tokens_when_none(): mock_message.content.text = "Test question" params.messages = [mock_message] params.temperature = None - params.maxTokens = None + params.maxTokens = 200 params.stopSequences = None params.systemPrompt = None params.tools = None @@ -1983,8 +1983,8 @@ async def test_mcp_tool_sampling_callback_omits_max_tokens_when_none(): assert isinstance(result, types.CreateMessageResult) call_kwargs = mock_chat_client.get_response.call_args - options = call_kwargs.kwargs.get("options") - assert options is None or "max_tokens" not in options + options = call_kwargs.kwargs.get("options") or {} + assert options["max_tokens"] == 200 async def test_connect_sampling_capabilities_with_client():