Python: Return MCP tool-use sampling results

python/packages/core/agent_framework/_mcp.py

@@ -780,13 +780,28 @@ def _prepare_content_for_mcp(
         self,
         content: Content,
     ) -> (
-        types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink | None
+        types.TextContent
+        | types.ImageContent
+        | types.AudioContent
+        | types.EmbeddedResource
+        | types.ResourceLink
+        | types.ToolUseContent
+        | None
     ):
         """Prepare an Agent Framework content type for MCP."""
         from mcp import types
 
         if content.type == "text":
             return types.TextContent(type="text", text=content.text)  # type: ignore[attr-defined]
+        if content.type == "function_call":
+            if not content.call_id or not content.name:
+                return None
+            return types.ToolUseContent(
+                type="tool_use",
+                id=content.call_id,
+                name=content.name,
+                input=content.parse_arguments() or {},
+            )
         if content.type == "data":
             if content.media_type and content.media_type.startswith("image/"):
                 return types.ImageContent(type="image", data=content.uri, mimeType=content.media_type)  # type: ignore[attr-defined]
@@ -822,11 +837,21 @@ def _prepare_message_for_mcp(
         self,
         content: Message,
     ) -> list[
-        types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink
+        types.TextContent
+        | types.ImageContent
+        | types.AudioContent
+        | types.EmbeddedResource
+        | types.ResourceLink
+        | types.ToolUseContent
     ]:
         """Prepare a Message for MCP format."""
         messages: list[
-            types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink
+            types.TextContent
+            | types.ImageContent
+            | types.AudioContent
+            | types.EmbeddedResource
+            | types.ResourceLink
+            | types.ToolUseContent
         ] = []
         for item in content.contents:
             mcp_content = self._prepare_content_for_mcp(item)
@@ -1433,7 +1458,7 @@ async def sampling_callback(
         self,
         context: RequestContext[ClientSession, Any],
         params: types.CreateMessageRequestParams,
-    ) -> types.CreateMessageResult | types.ErrorData:
+    ) -> types.CreateMessageResult | types.CreateMessageResultWithTools | types.ErrorData:
         """Callback function for sampling.
 
         This function is called when the MCP server sends a ``sampling/createMessage``
@@ -1460,8 +1485,8 @@ async def sampling_callback(
             params: The message creation request parameters.
 
         Returns:
-            Either a CreateMessageResult with the generated message or ErrorData if the request
-            is denied, rate limited, or generation fails.
+            Either a CreateMessageResult/CreateMessageResultWithTools with the generated message or ErrorData if the
+            request is denied, rate limited, or generation fails.
         """
         from mcp import types
 
@@ -1544,7 +1569,21 @@ async def sampling_callback(
                 code=types.INTERNAL_ERROR,
                 message="Failed to get chat message content.",
             )
-        mcp_contents = self._prepare_message_for_mcp(response.messages[0])
+        mcp_contents = [
+            mcp_content for message in response.messages for mcp_content in self._prepare_message_for_mcp(message)
+        ]
+        tool_use_contents: list[types.SamplingMessageContentBlock] = []
+        for content in mcp_contents:
+            if isinstance(content, types.ToolUseContent):
+                tool_use_contents.append(content)
+        if tool_use_contents:
+            return types.CreateMessageResultWithTools(
+                role="assistant",
+                content=tool_use_contents,
+                model=response.model or "unknown",
+                stopReason="toolUse",
+            )
+
         # grab the first content that is of type TextContent or ImageContent
         mcp_content = next(
             (content for content in mcp_contents if isinstance(content, (types.TextContent, types.ImageContent))),

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

@@ -2954,6 +2954,65 @@ async def test_mcp_tool_sampling_callback_no_response_and_successful_message_cre
     assert success.content.text == "Hello"
 
 
+async def test_mcp_tool_sampling_callback_returns_tool_use_results():
+    """Test sampling callback returns structured MCP tool-use content from function calls."""
+    tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
+
+    mock_chat_client = AsyncMock()
+    mock_chat_client.get_response.return_value = Mock(
+        messages=[
+            Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(
+                        call_id="call-1",
+                        name="Answer",
+                        arguments={"answer": "hello"},
+                    ),
+                ],
+            ),
+            Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(
+                        call_id="call-2",
+                        name="Citations",
+                        arguments={"sources": ["docs"]},
+                    ),
+                ],
+            ),
+        ],
+        model="test-model",
+    )
+    tool.client = mock_chat_client
+
+    params = Mock()
+    params.messages = [types.PromptMessage(role="user", content=types.TextContent(type="text", text="Answer"))]
+    params.temperature = None
+    params.maxTokens = 100
+    params.stopSequences = None
+    params.systemPrompt = None
+    params.tools = [
+        types.Tool(name="Answer", description="Return an answer", inputSchema={"type": "object"}),
+        types.Tool(name="Citations", description="Return source citations", inputSchema={"type": "object"}),
+    ]
+    params.toolChoice = None
+
+    result = await tool.sampling_callback(Mock(), params)
+
+    assert isinstance(result, types.CreateMessageResultWithTools)
+    assert result.role == "assistant"
+    assert result.model == "test-model"
+    assert result.stopReason == "toolUse"
+    assert isinstance(result.content, list)
+    tool_use_contents = [content for content in result.content if isinstance(content, types.ToolUseContent)]
+    assert tool_use_contents == result.content
+    assert [(content.id, content.name, content.input) for content in tool_use_contents] == [
+        ("call-1", "Answer", {"answer": "hello"}),
+        ("call-2", "Citations", {"sources": ["docs"]}),
+    ]
+
+
 async def test_mcp_tool_logging_callback_logs_at_requested_level() -> None:
     tool = MCPStdioTool(name="test_tool", command="python")