Skip to content

Commit d69d94e

Browse files
committed
Test and example MCP server to illustrate Issue modelcontextprotocol#201.
File mcp_stdio_client.py is adapted from the example in simple-chatbot.
1 parent f10665d commit d69d94e

File tree

3 files changed

+252
-0
lines changed

3 files changed

+252
-0
lines changed

tests/example_mcp_server.py

+33
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
#
2+
# Small Demo server using FastMCP and illustrating debugging and notification streams
3+
#
4+
5+
import logging
6+
from mcp.server.fastmcp import FastMCP, Context
7+
import time
8+
import asyncio
9+
10+
mcp = FastMCP("MCP EXAMPLE SERVER", debug=True, log_level="DEBUG")
11+
12+
logger = logging.getLogger(__name__)
13+
14+
logger.debug(f"MCP STARTING EXAMPLE SERVER")
15+
16+
@mcp.resource("config://app")
17+
def get_config() -> str:
18+
"""Static configuration data"""
19+
return "Test Server 2024-02-25"
20+
21+
@mcp.tool()
22+
async def simple_tool(x:float, y:float, ctx:Context) -> str:
23+
logger.debug("IN SIMPLE_TOOL")
24+
await ctx.report_progress(1, 2)
25+
return x*y
26+
27+
@mcp.tool()
28+
async def simple_tool_with_logging(x:float, y:float, ctx:Context) -> str:
29+
await ctx.info(f"Processing Simple Tool")
30+
logger.debug("IN SIMPLE_TOOL")
31+
await ctx.report_progress(1, 2)
32+
return x*y
33+

tests/mcp_stdio_client.py

+136
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
from mcp import ClientSession, ListToolsResult, StdioServerParameters
2+
from mcp.client.stdio import stdio_client
3+
from mcp.types import CallToolResult
4+
from mcp import Tool as MCPTool
5+
6+
from contextlib import AsyncExitStack
7+
from typing import Any
8+
import asyncio
9+
10+
11+
import logging
12+
logger = logging.getLogger(__name__)
13+
14+
15+
16+
class NotificationLoggingClientSession(ClientSession):
17+
18+
def __init__(self, read_stream, write_stream):
19+
print(f"NOTIFICATION LOGGING CLIENT SESSION")
20+
super().__init__(read_stream, write_stream)
21+
22+
# override base session to log incoming notifications
23+
async def _received_notification(self, notification):
24+
print(f"NOTIFICATION:{notification}")
25+
print(f"NOTIFICATION-END")
26+
27+
async def send_progress_notification(self, progress_token, progress, total):
28+
print(f"PROGRESS:{progress_token}")
29+
print(f"PROGRESS-END")
30+
31+
32+
# adapted from mcp-python-sdk/examples/clients/simple-chatbot/mcp_simple_chatbot/main.py
33+
class MCPClient:
34+
"""Manages MCP server connections and tool execution."""
35+
36+
def __init__(self, name, server_params: StdioServerParameters, errlog=None):
37+
self.name = name
38+
self.server_params = server_params
39+
self.errlog = errlog
40+
self.stdio_context: Any | None = None
41+
self.session: ClientSession | None = None
42+
self._cleanup_lock: asyncio.Lock = asyncio.Lock()
43+
self.exit_stack: AsyncExitStack = AsyncExitStack()
44+
45+
async def initialize(self) -> None:
46+
"""Initialize the server connection."""
47+
48+
try:
49+
stdio_transport = await self.exit_stack.enter_async_context(
50+
stdio_client(self.server_params)
51+
)
52+
read, write = stdio_transport
53+
session = await self.exit_stack.enter_async_context(
54+
# ClientSession(read, write)
55+
NotificationLoggingClientSession(read, write)
56+
)
57+
await session.initialize()
58+
self.session = session
59+
except Exception as e:
60+
logging.error(f"Error initializing server: {e}")
61+
await self.cleanup()
62+
raise
63+
64+
async def get_available_tools(self) -> list[MCPTool]:
65+
"""List available tools from the server.
66+
67+
Returns:
68+
A list of available tools.
69+
70+
Raises:
71+
RuntimeError: If the server is not initialized.
72+
"""
73+
if not self.session:
74+
raise RuntimeError(f"Server {self.name} not initialized")
75+
76+
tools_response = await self.session.list_tools()
77+
78+
# Let's just ignore pagination for now
79+
return tools_response.tools
80+
81+
async def call_tool(
82+
self,
83+
tool_name: str,
84+
arguments: dict[str, Any],
85+
retries: int = 2,
86+
delay: float = 1.0,
87+
) -> Any:
88+
"""Execute a tool with retry mechanism.
89+
90+
Args:
91+
tool_name: Name of the tool to execute.
92+
arguments: Tool arguments.
93+
retries: Number of retry attempts.
94+
delay: Delay between retries in seconds.
95+
96+
Returns:
97+
Tool execution result.
98+
99+
Raises:
100+
RuntimeError: If server is not initialized.
101+
Exception: If tool execution fails after all retries.
102+
"""
103+
if not self.session:
104+
raise RuntimeError(f"Server {self.name} not initialized")
105+
106+
attempt = 0
107+
while attempt < retries:
108+
try:
109+
logging.info(f"Executing {tool_name}...")
110+
result = await self.session.call_tool(tool_name, arguments)
111+
112+
return result
113+
114+
except Exception as e:
115+
attempt += 1
116+
logging.warning(
117+
f"Error executing tool: {e}. Attempt {attempt} of {retries}."
118+
)
119+
if attempt < retries:
120+
logging.info(f"Retrying in {delay} seconds...")
121+
await asyncio.sleep(delay)
122+
else:
123+
logging.error("Max retries reached. Failing.")
124+
raise
125+
126+
async def cleanup(self) -> None:
127+
"""Clean up server resources."""
128+
async with self._cleanup_lock:
129+
try:
130+
await self.exit_stack.aclose()
131+
self.session = None
132+
self.stdio_context = None
133+
except Exception as e:
134+
logging.error(f"Error during cleanup of server {self.name}: {e}")
135+
136+

tests/test_mcp_tool.py

+83
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
import pytest
2+
3+
import os
4+
import sys
5+
from .mcp_stdio_client import MCPClient
6+
7+
from mcp import StdioServerParameters
8+
9+
# locate the exmaple MCP server co-located in this directory
10+
11+
mcp_server_dir = os.path.dirname(os.path.abspath(__file__))
12+
mcp_server_file = os.path.join(mcp_server_dir, "example_mcp_server.py")
13+
14+
# mcpServers config in same syntax used by reference MCP
15+
16+
servers_config = {
17+
"mcpServers": {
18+
19+
"testMcpServer": {
20+
"command": "mcp", # be sure to . .venv/bin/activate so that mcp command is found
21+
"args": [
22+
"run",
23+
mcp_server_file
24+
]
25+
}
26+
27+
}
28+
}
29+
30+
31+
# @pytest.mark.asyncio
32+
@pytest.mark.anyio
33+
async def test_mcp():
34+
35+
servers = servers_config.get("mcpServers")
36+
37+
server0 = "testMcpServer"
38+
config0 = servers[server0]
39+
40+
client = MCPClient(
41+
server0,
42+
StdioServerParameters.model_validate(config0)
43+
)
44+
await client.initialize()
45+
tools = await client.get_available_tools()
46+
47+
print(f"TOOLS:{tools}")
48+
mcp_tool = tools[0]
49+
50+
res = await client.call_tool("simple_tool", {"x":5, "y":7})
51+
52+
print(f"RES:{res}")
53+
54+
# clients must be destroyed in reverse order
55+
await client.cleanup()
56+
57+
58+
# @pytest.mark.asyncio
59+
@pytest.mark.anyio
60+
async def test_mcp_with_logging():
61+
62+
servers = servers_config.get("mcpServers")
63+
64+
server0 = "testMcpServer"
65+
config0 = servers[server0]
66+
67+
client = MCPClient(
68+
server0,
69+
StdioServerParameters.model_validate(config0)
70+
)
71+
await client.initialize()
72+
tools = await client.get_available_tools()
73+
74+
print(f"TOOLS:{tools}")
75+
mcp_tool = tools[0]
76+
77+
res = await client.call_tool("simple_tool_with_logging", {"x":5, "y":7})
78+
79+
print(f"RES:{res}")
80+
81+
# clients must be destroyed in reverse order
82+
await client.cleanup()
83+

0 commit comments

Comments
 (0)