Python: Introducing Local MCP Servers (#389)

* mcp parts * mcp parts 2 * removed structured output in favor of handling in chatresponse, mcp as AITool and running samples * updated naming * fixed test
2026-06-16 21:04:09 +08:00 · 2025-08-13 11:48:22 +02:00
parent 80b0920e58
commit ad3d8171bf
20 changed files with 1970 additions and 298 deletions
@@ -0,0 +1,544 @@
+# Copyright (c) Microsoft. All rights reserved.
+# type: ignore[reportPrivateUsage]
+import os
+from contextlib import _AsyncGeneratorContextManager  # type: ignore
+from typing import Any
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+from mcp import types
+from mcp.client.session import ClientSession
+from mcp.shared.exceptions import McpError
+from pydantic import AnyUrl, ValidationError
+
+from agent_framework import (
+    AITool,
+    ChatMessage,
+    ChatRole,
+    DataContent,
+    McpSseTools,
+    McpStdioTool,
+    McpStreamableHttpTool,
+    McpWebsocketTool,
+    TextContent,
+    UriContent,
+)
+from agent_framework._mcp import (
+    McpTool,
+    _ai_content_to_mcp_types,
+    _chat_message_to_mcp_types,
+    _get_input_model_from_mcp_prompt,
+    _get_input_model_from_mcp_tool,
+    _mcp_call_tool_result_to_ai_contents,
+    _mcp_prompt_message_to_chat_message,
+    _mcp_type_to_ai_content,
+    _normalize_mcp_name,
+)
+from agent_framework.exceptions import ToolExecutionException
+
+# Integration test skip condition
+skip_if_mcp_integration_tests_disabled = pytest.mark.skipif(
+    os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true" or os.getenv("LOCAL_MCP_URL", "") == "",
+    reason="No LOCAL_MCP_URL provided; skipping integration tests."
+    if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
+    else "Integration tests are disabled.",
+)
+
+
+# Helper function tests
+def test_normalize_mcp_name():
+    """Test MCP name normalization."""
+    assert _normalize_mcp_name("valid_name") == "valid_name"
+    assert _normalize_mcp_name("name-with-dashes") == "name-with-dashes"
+    assert _normalize_mcp_name("name.with.dots") == "name.with.dots"
+    assert _normalize_mcp_name("name with spaces") == "name-with-spaces"
+    assert _normalize_mcp_name("name@with#special$chars") == "name-with-special-chars"
+    assert _normalize_mcp_name("name/with\\slashes") == "name-with-slashes"
+
+
+def test_mcp_prompt_message_to_ai_content():
+    """Test conversion from MCP prompt message to AI content."""
+    mcp_message = types.PromptMessage(role="user", content=types.TextContent(type="text", text="Hello, world!"))
+    ai_content = _mcp_prompt_message_to_chat_message(mcp_message)
+
+    assert isinstance(ai_content, ChatMessage)
+    assert ai_content.role.value == "user"
+    assert len(ai_content.contents) == 1
+    assert isinstance(ai_content.contents[0], TextContent)
+    assert ai_content.contents[0].text == "Hello, world!"
+    assert ai_content.raw_representation == mcp_message
+
+
+def test_mcp_call_tool_result_to_ai_contents():
+    """Test conversion from MCP tool result to AI contents."""
+    mcp_result = types.CallToolResult(
+        content=[
+            types.TextContent(type="text", text="Result text"),
+            types.ImageContent(type="image", data="data:image/png;base64,xyz", mimeType="image/png"),
+        ]
+    )
+    ai_contents = _mcp_call_tool_result_to_ai_contents(mcp_result)
+
+    assert len(ai_contents) == 2
+    assert isinstance(ai_contents[0], TextContent)
+    assert ai_contents[0].text == "Result text"
+    assert isinstance(ai_contents[1], DataContent)
+    assert ai_contents[1].uri == "data:image/png;base64,xyz"
+    assert ai_contents[1].media_type == "image/png"
+
+
+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")
+    ai_content = _mcp_type_to_ai_content(mcp_content)
+
+    assert isinstance(ai_content, TextContent)
+    assert ai_content.text == "Sample text"
+    assert ai_content.raw_representation == mcp_content
+
+
+def test_mcp_content_types_to_ai_content_image():
+    """Test conversion of MCP image content to AI content."""
+    mcp_content = types.ImageContent(type="image", data="data:image/jpeg;base64,abc", mimeType="image/jpeg")
+    ai_content = _mcp_type_to_ai_content(mcp_content)
+
+    assert isinstance(ai_content, DataContent)
+    assert ai_content.uri == "data:image/jpeg;base64,abc"
+    assert ai_content.media_type == "image/jpeg"
+    assert ai_content.raw_representation == mcp_content
+
+
+def test_mcp_content_types_to_ai_content_audio():
+    """Test conversion of MCP audio content to AI content."""
+    mcp_content = types.AudioContent(type="audio", data="data:audio/wav;base64,def", mimeType="audio/wav")
+    ai_content = _mcp_type_to_ai_content(mcp_content)
+
+    assert isinstance(ai_content, DataContent)
+    assert ai_content.uri == "data:audio/wav;base64,def"
+    assert ai_content.media_type == "audio/wav"
+    assert ai_content.raw_representation == mcp_content
+
+
+def test_mcp_content_types_to_ai_content_resource_link():
+    """Test conversion of MCP resource link to AI content."""
+    mcp_content = types.ResourceLink(
+        type="resource_link",
+        uri=AnyUrl("https://example.com/resource"),
+        name="test_resource",
+        mimeType="application/json",
+    )
+    ai_content = _mcp_type_to_ai_content(mcp_content)
+
+    assert isinstance(ai_content, UriContent)
+    assert ai_content.uri == "https://example.com/resource"
+    assert ai_content.media_type == "application/json"
+    assert ai_content.raw_representation == mcp_content
+
+
+def test_mcp_content_types_to_ai_content_embedded_resource_text():
+    """Test conversion of MCP embedded text resource to AI content."""
+    text_resource = types.TextResourceContents(
+        uri=AnyUrl("file://test.txt"), mimeType="text/plain", text="Embedded text content"
+    )
+    mcp_content = types.EmbeddedResource(type="resource", resource=text_resource)
+    ai_content = _mcp_type_to_ai_content(mcp_content)
+
+    assert isinstance(ai_content, TextContent)
+    assert ai_content.text == "Embedded text content"
+    assert ai_content.raw_representation == mcp_content
+
+
+def test_mcp_content_types_to_ai_content_embedded_resource_blob():
+    """Test conversion of MCP embedded blob resource to AI content."""
+    # Use a proper data URI in the blob field since that's what the MCP implementation expects
+    blob_resource = types.BlobResourceContents(
+        uri=AnyUrl("file://test.bin"),
+        mimeType="application/octet-stream",
+        blob="data:application/octet-stream;base64,dGVzdCBkYXRh",
+    )
+    mcp_content = types.EmbeddedResource(type="resource", resource=blob_resource)
+    ai_content = _mcp_type_to_ai_content(mcp_content)
+
+    assert isinstance(ai_content, DataContent)
+    assert ai_content.uri == "data:application/octet-stream;base64,dGVzdCBkYXRh"
+    assert ai_content.media_type == "application/octet-stream"
+    assert ai_content.raw_representation == mcp_content
+
+
+def test_ai_content_to_mcp_content_types_text():
+    """Test conversion of AI text content to MCP content."""
+    ai_content = TextContent(text="Sample text")
+    mcp_content = _ai_content_to_mcp_types(ai_content)
+
+    assert isinstance(mcp_content, types.TextContent)
+    assert mcp_content.type == "text"
+    assert mcp_content.text == "Sample text"
+
+
+def test_ai_content_to_mcp_content_types_data_image():
+    """Test conversion of AI data content to MCP content."""
+    ai_content = DataContent(uri="data:image/png;base64,xyz", media_type="image/png")
+    mcp_content = _ai_content_to_mcp_types(ai_content)
+
+    assert isinstance(mcp_content, types.ImageContent)
+    assert mcp_content.type == "image"
+    assert mcp_content.data == "data:image/png;base64,xyz"
+    assert mcp_content.mimeType == "image/png"
+
+
+def test_ai_content_to_mcp_content_types_data_audio():
+    """Test conversion of AI data content to MCP content."""
+    ai_content = DataContent(uri="data:audio/mpeg;base64,xyz", media_type="audio/mpeg")
+    mcp_content = _ai_content_to_mcp_types(ai_content)
+
+    assert isinstance(mcp_content, types.AudioContent)
+    assert mcp_content.type == "audio"
+    assert mcp_content.data == "data:audio/mpeg;base64,xyz"
+    assert mcp_content.mimeType == "audio/mpeg"
+
+
+def test_ai_content_to_mcp_content_types_data_binary():
+    """Test conversion of AI data content to MCP content."""
+    ai_content = DataContent(uri="data:application/octet-stream;base64,xyz", media_type="application/octet-stream")
+    mcp_content = _ai_content_to_mcp_types(ai_content)
+
+    assert isinstance(mcp_content, types.EmbeddedResource)
+    assert mcp_content.type == "resource"
+    assert mcp_content.resource.blob == "data:application/octet-stream;base64,xyz"
+    assert mcp_content.resource.mimeType == "application/octet-stream"
+
+
+def test_ai_content_to_mcp_content_types_uri():
+    """Test conversion of AI URI content to MCP content."""
+    ai_content = UriContent(uri="https://example.com/resource", media_type="application/json")
+    mcp_content = _ai_content_to_mcp_types(ai_content)
+
+    assert isinstance(mcp_content, types.ResourceLink)
+    assert mcp_content.type == "resource_link"
+    assert str(mcp_content.uri) == "https://example.com/resource"
+    assert mcp_content.mimeType == "application/json"
+
+
+def test_chat_message_to_mcp_types():
+    message = ChatMessage(
+        role="user",
+        contents=[TextContent(text="test"), DataContent(uri="data:image/png;base64,xyz", media_type="image/png")],
+    )
+    mcp_contents = _chat_message_to_mcp_types(message)
+    assert len(mcp_contents) == 2
+    assert isinstance(mcp_contents[0], types.TextContent)
+    assert isinstance(mcp_contents[1], types.ImageContent)
+
+
+def test_get_input_model_from_mcp_tool():
+    """Test creation of input model from MCP tool."""
+    tool = types.Tool(
+        name="test_tool",
+        description="A test tool",
+        inputSchema={
+            "type": "object",
+            "properties": {"param1": {"type": "string"}, "param2": {"type": "number"}},
+            "required": ["param1"],
+        },
+    )
+    model = _get_input_model_from_mcp_tool(tool)
+
+    # Create an instance to verify the model works
+    instance = model(param1="test", param2=42)
+    assert instance.param1 == "test"
+    assert instance.param2 == 42
+
+    # Test validation
+    with pytest.raises(ValidationError):  # Missing required param1
+        model(param2=42)
+
+
+def test_get_input_model_from_mcp_prompt():
+    """Test creation of input model from MCP prompt."""
+    prompt = types.Prompt(
+        name="test_prompt",
+        description="A test prompt",
+        arguments=[
+            types.PromptArgument(name="arg1", description="First argument", required=True),
+            types.PromptArgument(name="arg2", description="Second argument", required=False),
+        ],
+    )
+    model = _get_input_model_from_mcp_prompt(prompt)
+
+    # Create an instance to verify the model works
+    instance = model(arg1="test", arg2="optional")
+    assert instance.arg1 == "test"
+    assert instance.arg2 == "optional"
+
+    # Test validation
+    with pytest.raises(ValidationError):  # Missing required arg1
+        model(arg2="optional")
+
+
+# McpTool tests
+async def test_local_mcp_server_initialization():
+    """Test McpTool initialization."""
+    server = McpTool(name="test_server")
+    assert isinstance(server, AITool)
+    assert server.name == "test_server"
+    assert server.session is None
+    assert server.functions == []
+
+
+async def test_local_mcp_server_context_manager():
+    """Test McpTool as context manager."""
+
+    class TestServer(McpTool):
+        async def connect(self):
+            # Mock connection
+            self.session = Mock(spec=ClientSession)
+
+        def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+            return None
+
+    server = TestServer(name="test_server")
+    async with server:
+        assert server.session is not None
+
+    assert server.session is None
+
+
+async def test_local_mcp_server_load_functions():
+    """Test loading functions from MCP server."""
+
+    class TestServer(McpTool):
+        async def connect(self):
+            self.session = Mock(spec=ClientSession)
+            # Mock tools list response
+            self.session.list_tools = AsyncMock(
+                return_value=types.ListToolsResult(
+                    tools=[
+                        types.Tool(
+                            name="test_tool",
+                            description="Test tool",
+                            inputSchema={
+                                "type": "object",
+                                "properties": {"param": {"type": "string"}},
+                                "required": ["param"],
+                            },
+                        )
+                    ]
+                )
+            )
+
+        def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+            return None
+
+    server = TestServer(name="test_server")
+    assert isinstance(server, AITool)
+    async with server:
+        await server.load_tools()
+        assert len(server.functions) == 1
+        assert server.functions[0].name == "test_tool"
+
+
+async def test_local_mcp_server_load_prompts():
+    """Test loading prompts from MCP server."""
+
+    class TestServer(McpTool):
+        async def connect(self):
+            self.session = Mock(spec=ClientSession)
+            # Mock prompts list response
+            self.session.list_prompts = AsyncMock(
+                return_value=types.ListPromptsResult(
+                    prompts=[
+                        types.Prompt(
+                            name="test_prompt",
+                            description="Test prompt",
+                            arguments=[types.PromptArgument(name="arg", description="Test arg", required=True)],
+                        )
+                    ]
+                )
+            )
+
+        def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+            return None
+
+    server = TestServer(name="test_server")
+    async with server:
+        await server.load_prompts()
+        assert len(server.functions) == 1
+        assert server.functions[0].name == "test_prompt"
+
+
+async def test_local_mcp_server_function_execution():
+    """Test function execution through MCP server."""
+
+    class TestServer(McpTool):
+        async def connect(self):
+            self.session = Mock(spec=ClientSession)
+            self.session.list_tools = AsyncMock(
+                return_value=types.ListToolsResult(
+                    tools=[
+                        types.Tool(
+                            name="test_tool",
+                            description="Test tool",
+                            inputSchema={
+                                "type": "object",
+                                "properties": {"param": {"type": "string"}},
+                                "required": ["param"],
+                            },
+                        )
+                    ]
+                )
+            )
+            self.session.call_tool = AsyncMock(
+                return_value=types.CallToolResult(
+                    content=[types.TextContent(type="text", text="Tool executed successfully")]
+                )
+            )
+
+        def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+            return None
+
+    server = TestServer(name="test_server")
+    async with server:
+        await server.load_tools()
+        func = server.functions[0]
+        result = await func.invoke(param="test_value")
+
+        assert len(result) == 1
+        assert isinstance(result[0], TextContent)
+        assert result[0].text == "Tool executed successfully"
+
+
+async def test_local_mcp_server_function_execution_error():
+    """Test function execution error handling."""
+
+    class TestServer(McpTool):
+        async def connect(self):
+            self.session = Mock(spec=ClientSession)
+            self.session.list_tools = AsyncMock(
+                return_value=types.ListToolsResult(
+                    tools=[
+                        types.Tool(
+                            name="test_tool",
+                            description="Test tool",
+                            inputSchema={
+                                "type": "object",
+                                "properties": {"param": {"type": "string"}},
+                                "required": ["param"],
+                            },
+                        )
+                    ]
+                )
+            )
+            # Mock a tool call that raises an MCP error
+            self.session.call_tool = AsyncMock(
+                side_effect=McpError(types.ErrorData(code=-1, message="Tool execution failed"))
+            )
+
+        def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+            return None
+
+    server = TestServer(name="test_server")
+    async with server:
+        await server.load_tools()
+        func = server.functions[0]
+
+        with pytest.raises(ToolExecutionException):
+            await func.invoke(param="test_value")
+
+
+async def test_local_mcp_server_prompt_execution():
+    """Test prompt execution through MCP server."""
+
+    class TestMcpTool(McpTool):
+        async def connect(self):
+            self.session = Mock(spec=ClientSession)
+            self.session.list_prompts = AsyncMock(
+                return_value=types.ListPromptsResult(
+                    prompts=[
+                        types.Prompt(
+                            name="test_prompt",
+                            description="Test prompt",
+                            arguments=[types.PromptArgument(name="arg", description="Test arg", required=True)],
+                        )
+                    ]
+                )
+            )
+            self.session.get_prompt = AsyncMock(
+                return_value=types.GetPromptResult(
+                    description="Generated prompt",
+                    messages=[
+                        types.PromptMessage(role="user", content=types.TextContent(type="text", text="Test message"))
+                    ],
+                )
+            )
+
+        def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+            return None
+
+    server = TestMcpTool(name="test_server")
+    async with server:
+        await server.load_prompts()
+        prompt = server.functions[0]
+        result = await prompt.invoke(arg="test_value")
+
+        assert len(result) == 1
+        assert isinstance(result[0], ChatMessage)
+        assert result[0].role == ChatRole.USER
+        assert len(result[0].contents) == 1
+        assert result[0].contents[0].text == "Test message"
+
+
+# Server implementation tests
+def test_local_mcp_stdio_tool_init():
+    """Test McpStdioTool initialization."""
+    tool = McpStdioTool(name="test", command="echo", args=["hello"])
+    assert tool.name == "test"
+    assert tool.command == "echo"
+    assert tool.args == ["hello"]
+
+
+def test_local_mcp_sse_tools_init():
+    """Test McpSseTools initialization."""
+    tool = McpSseTools(name="test", url="http://localhost:8080")
+    assert tool.name == "test"
+    assert tool.url == "http://localhost:8080"
+
+
+def test_local_mcp_websocket_tool_init():
+    """Test McpWebsocketTool initialization."""
+    tool = McpWebsocketTool(name="test", url="ws://localhost:8080")
+    assert tool.name == "test"
+    assert tool.url == "ws://localhost:8080"
+
+
+def test_local_mcp_streamable_http_tool_init():
+    """Test McpStreamableHttpTool initialization."""
+    tool = McpStreamableHttpTool(name="test", url="http://localhost:8080")
+    assert tool.name == "test"
+    assert tool.url == "http://localhost:8080"
+
+
+# Integration test
+@skip_if_mcp_integration_tests_disabled
+async def test_streamable_http_integration():
+    """Test MCP StreamableHTTP integration."""
+    url = os.environ.get("LOCAL_MCP_URL", "")
+    if not url.startswith("http"):
+        pytest.skip("LOCAL_MCP_URL is not an HTTP URL")
+
+    tool = McpStreamableHttpTool(name="integration_test", url=url)
+
+    async with tool:
+        # Test that we can connect and load tools
+        assert tool.session is not None
+        assert isinstance(tool.functions, list)
+
+        # If there are functions available, try to get information about one
+        assert tool.functions, "The MCP server should have at least one function."
+
+        func = tool.functions[0]
+
+        assert hasattr(func, "name")
+        assert hasattr(func, "description")
+
+        result = await func.invoke(query="What is Agent Framework?")
+        assert result[0].text is not None
@@ -31,7 +31,6 @@ from agent_framework import (
    HostedFileContent,
    HostedVectorStoreContent,
    SpeechToTextOptions,
-    StructuredResponse,
    TextContent,
    TextReasoningContent,
    TextSpanRegion,
@@ -472,27 +471,44 @@ def test_chat_response():
    assert str(response) == response.text


-# region StructuredResponse
+class OutputModel(BaseModel):
+    response: str


-def test_structured_response():
-    """Test the StructuredResponse class to ensure it initializes correctly with a value."""
+def test_chat_response_with_format():
+    """Test the ChatResponse class to ensure it initializes correctly with a message."""
+    # Create a ChatMessage
+    message = ChatMessage(role="assistant", text='{"response": "Hello"}')

-    class ResponseModel(BaseModel):
-        content: str
-        action: str
-
-    # Create a StructuredResponse with a value
-    response = StructuredResponse[ResponseModel](
-        value=ResponseModel(content="Hello, world!", action="test"),
-        text="{'content': 'Hello, world!', 'action': 'test'}",
-    )
+    # Create a ChatResponse with the message
+    response = ChatResponse(messages=message)

    # Check the type and content
-    assert response.value == ResponseModel(content="Hello, world!", action="test")
-    assert isinstance(response, StructuredResponse)
-    # text property returns joined messages text (single message present)
-    assert isinstance(response.text, str)
+    assert response.messages[0].role == ChatRole.ASSISTANT
+    assert response.messages[0].text == '{"response": "Hello"}'
+    assert isinstance(response.messages[0], ChatMessage)
+    assert response.text == '{"response": "Hello"}'
+    assert response.value is None
+    response.try_parse_value(OutputModel)
+    assert response.value is not None
+    assert response.value.response == "Hello"
+
+
+def test_chat_response_with_format_init():
+    """Test the ChatResponse class to ensure it initializes correctly with a message."""
+    # Create a ChatMessage
+    message = ChatMessage(role="assistant", text='{"response": "Hello"}')
+
+    # Create a ChatResponse with the message
+    response = ChatResponse(messages=message, response_format=OutputModel)
+
+    # Check the type and content
+    assert response.messages[0].role == ChatRole.ASSISTANT
+    assert response.messages[0].text == '{"response": "Hello"}'
+    assert isinstance(response.messages[0], ChatMessage)
+    assert response.text == '{"response": "Hello"}'
+    assert response.value is not None
+    assert response.value.response == "Hello"


 # region ChatResponseUpdate
@@ -636,6 +652,32 @@ async def test_chat_response_from_async_generator():
    assert resp.text == "Hello world"


+@mark.asyncio
+async def test_chat_response_from_async_generator_output_format():
+    async def gen() -> AsyncIterable[ChatResponseUpdate]:
+        yield ChatResponseUpdate(text='{ "respon', message_id="1")
+        yield ChatResponseUpdate(text='se": "Hello" }', message_id="1")
+
+    resp = await ChatResponse.from_chat_response_generator(gen())
+    assert resp.text == '{ "response": "Hello" }'
+    assert resp.value is None
+    resp.try_parse_value(OutputModel)
+    assert resp.value is not None
+    assert resp.value.response == "Hello"
+
+
+@mark.asyncio
+async def test_chat_response_from_async_generator_output_format_in_method():
+    async def gen() -> AsyncIterable[ChatResponseUpdate]:
+        yield ChatResponseUpdate(text='{ "respon', message_id="1")
+        yield ChatResponseUpdate(text='se": "Hello" }', message_id="1")
+
+    resp = await ChatResponse.from_chat_response_generator(gen(), output_format_type=OutputModel)
+    assert resp.text == '{ "response": "Hello" }'
+    assert resp.value is not None
+    assert resp.value.response == "Hello"
+
+
 # region ChatToolMode