Skip to content

Commit

Permalink
refactor: update test files to use ParsedMessage
Browse files Browse the repository at this point in the history
Updates test files to work with the ParsedMessage stream type aliases
and fixes a line length issue in test_201_client_hangs_on_logging.py.

Github-Issue:#201
  • Loading branch information
dsp-ant committed Mar 3, 2025
1 parent 8c9285e commit d83f811
Show file tree
Hide file tree
Showing 4 changed files with 117 additions and 29 deletions.
29 changes: 16 additions & 13 deletions tests/client/test_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import pytest

from mcp.client.session import ClientSession
from mcp.shared.session import ParsedMessage
from mcp.types import (
LATEST_PROTOCOL_VERSION,
ClientNotification,
Expand All @@ -11,7 +12,6 @@
InitializeRequest,
InitializeResult,
JSONRPCMessage,
JSONRPCNotification,
JSONRPCRequest,
JSONRPCResponse,
ServerCapabilities,
Expand All @@ -22,10 +22,10 @@
@pytest.mark.anyio
async def test_client_session_initialize():
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
JSONRPCMessage
ParsedMessage[None]
](1)
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
JSONRPCMessage
ParsedMessage[None]
](1)

initialized_notification = None
Expand Down Expand Up @@ -57,20 +57,23 @@ async def mock_server():

async with server_to_client_send:
await server_to_client_send.send(
JSONRPCMessage(
JSONRPCResponse(
jsonrpc="2.0",
id=jsonrpc_request.root.id,
result=result.model_dump(
by_alias=True, mode="json", exclude_none=True
),
)
ParsedMessage(
root=JSONRPCMessage(
JSONRPCResponse(
jsonrpc="2.0",
id=jsonrpc_request.root.id,
result=result.model_dump(
by_alias=True, mode="json", exclude_none=True
),
)
),
raw=None,
)
)
jsonrpc_notification = await client_to_server_receive.receive()
assert isinstance(jsonrpc_notification.root, JSONRPCNotification)
assert isinstance(jsonrpc_notification.root, ParsedMessage)
initialized_notification = ClientNotification.model_validate(
jsonrpc_notification.model_dump(
jsonrpc_notification.root.model_dump(
by_alias=True, mode="json", exclude_none=True
)
)
Expand Down
72 changes: 72 additions & 0 deletions tests/issues/test_201_client_hangs_on_logging.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
import logging

import anyio
import pytest

from mcp.server.fastmcp import Context, FastMCP
from mcp.shared.memory import create_connected_server_and_client_session
from mcp.types import TextContent

logger = logging.getLogger(__name__)


@pytest.mark.anyio
async def test_client_hangs_on_logging():
"""
Test case for issue #201: Client hangs when ctx.info() logging with pure Python
client.
This test creates a FastMCP server with a tool that uses ctx.info() for logging.
The issue is that when this tool is called, the client hangs and doesn't complete.
"""
# Create a FastMCP server with a tool that uses ctx.info() for logging
server = FastMCP("test-server")

@server.tool()
async def simple_tool(x: float, y: float) -> str:
"""A simple tool without logging"""
return str(x * y)

@server.tool()
async def tool_with_logging(x: float, y: float, ctx: Context) -> str:
"""A tool that uses ctx.info() for logging - this causes the client to hang"""
await ctx.info(f"Processing tool with x={x} and y={y}")
logger.debug("Inside tool_with_logging")
await ctx.report_progress(1, 2)
return str(x * y)

# Create a client session connected to the server
async with create_connected_server_and_client_session(
server._mcp_server
) as client_session:
# First test that a simple tool works correctly
result = await client_session.call_tool("simple_tool", {"x": 5, "y": 7})
assert isinstance(result.content[0], TextContent)
assert result.content[0].text == "35.0"

# Create an event to signal when the task is done
done_event = anyio.Event()

# Use a separate task for calling the tool that might hang
async def call_logging_tool():
try:
with anyio.fail_after(5): # Add a timeout to prevent test from hanging
result = await client_session.call_tool(
"tool_with_logging", {"x": 3, "y": 4}
)
assert isinstance(result.content[0], TextContent)
assert result.content[0].text == "12.0"
done_event.set()
except Exception as e:
logger.error(f"Error calling tool_with_logging: {e}")
raise

# Start the task
async with anyio.create_task_group() as tg:
tg.start_soon(call_logging_tool)

# Wait for the task to complete or timeout
with anyio.fail_after(10):
await done_event.wait()

# If we get here without hanging or timeout exceptions, the test passes
6 changes: 3 additions & 3 deletions tests/server/test_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,10 @@
from mcp.server.lowlevel import NotificationOptions
from mcp.server.models import InitializationOptions
from mcp.server.session import ServerSession
from mcp.shared.session import ParsedMessage
from mcp.types import (
ClientNotification,
InitializedNotification,
JSONRPCMessage,
PromptsCapability,
ResourcesCapability,
ServerCapabilities,
Expand All @@ -19,10 +19,10 @@
@pytest.mark.anyio
async def test_server_session_initialize():
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
JSONRPCMessage
ParsedMessage[None]
](1)
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
JSONRPCMessage
ParsedMessage[None]
](1)

async def run_client(client: ClientSession):
Expand Down
39 changes: 26 additions & 13 deletions tests/server/test_stdio.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import pytest

from mcp.server.stdio import stdio_server
from mcp.shared.session import ParsedMessage
from mcp.types import JSONRPCMessage, JSONRPCRequest, JSONRPCResponse


Expand All @@ -13,8 +14,12 @@ async def test_stdio_server():
stdout = io.StringIO()

messages = [
JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")),
JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=2, result={})),
ParsedMessage(
root=JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping"))
),
ParsedMessage(
root=JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=2, result={}))
),
]

for message in messages:
Expand All @@ -35,17 +40,25 @@ async def test_stdio_server():

# Verify received messages
assert len(received_messages) == 2
assert received_messages[0] == JSONRPCMessage(
root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
assert received_messages[0] == ParsedMessage(
root=JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping"))
)
assert received_messages[1] == JSONRPCMessage(
root=JSONRPCResponse(jsonrpc="2.0", id=2, result={})
assert received_messages[1] == ParsedMessage(
root=JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=2, result={}))
)

# Test sending responses from the server
responses = [
JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=3, method="ping")),
JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=4, result={})),
ParsedMessage(
root=JSONRPCMessage(
root=JSONRPCRequest(jsonrpc="2.0", id=3, method="ping")
)
),
ParsedMessage(
root=JSONRPCMessage(
root=JSONRPCResponse(jsonrpc="2.0", id=4, result={})
)
),
]

async with write_stream:
Expand All @@ -57,12 +70,12 @@ async def test_stdio_server():
assert len(output_lines) == 2

received_responses = [
JSONRPCMessage.model_validate_json(line.strip()) for line in output_lines
ParsedMessage.model_validate_json(line.strip()) for line in output_lines
]
assert len(received_responses) == 2
assert received_responses[0] == JSONRPCMessage(
root=JSONRPCRequest(jsonrpc="2.0", id=3, method="ping")
assert received_responses[0] == ParsedMessage(
root=JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=3, method="ping"))
)
assert received_responses[1] == JSONRPCMessage(
root=JSONRPCResponse(jsonrpc="2.0", id=4, result={})
assert received_responses[1] == ParsedMessage(
root=JSONRPCMessage(root=JSONRPCResponse(jsonrpc="2.0", id=4, result={}))
)

0 comments on commit d83f811

Please sign in to comment.