mirror of
https://github.com/aljazceru/mcp-python-sdk.git
synced 2025-12-19 06:54:18 +01:00
9d0f2daddb
* refactor: improve typing with memory stream type aliases Move memory stream type definitions to models.py and use them throughout the codebase for better type safety and maintainability. GitHub-Issue:#201 * refactor: move streams to ParsedMessage * refactor: update test files to use ParsedMessage 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 * refactor: rename ParsedMessage to MessageFrame for clarity 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com> * refactor: move MessageFrame class to types.py for better code organization 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com> * fix pyright * refactor: update websocket client to use MessageFrame Modified the websocket client to work with the new MessageFrame type, preserving raw message text and properly extracting the root JSON-RPC message when sending. Github-Issue:#204 * fix: use NoneType instead of None for type parameters in MessageFrame 🤖 Generated with [Claude Code](https://claude.ai/code) * refactor: rename root to message
111 lines
3.7 KiB
Python
111 lines
3.7 KiB
Python
from types import NoneType
|
|
|
|
import anyio
|
|
import pytest
|
|
|
|
from mcp.client.session import ClientSession
|
|
from mcp.types import (
|
|
LATEST_PROTOCOL_VERSION,
|
|
ClientNotification,
|
|
ClientRequest,
|
|
Implementation,
|
|
InitializedNotification,
|
|
InitializeRequest,
|
|
InitializeResult,
|
|
JSONRPCMessage,
|
|
JSONRPCRequest,
|
|
JSONRPCResponse,
|
|
MessageFrame,
|
|
ServerCapabilities,
|
|
ServerResult,
|
|
)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_client_session_initialize():
|
|
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
|
|
MessageFrame[NoneType]
|
|
](1)
|
|
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
|
|
MessageFrame[NoneType]
|
|
](1)
|
|
|
|
initialized_notification = None
|
|
|
|
async def mock_server():
|
|
nonlocal initialized_notification
|
|
|
|
jsonrpc_request = await client_to_server_receive.receive()
|
|
assert isinstance(jsonrpc_request, MessageFrame)
|
|
request = ClientRequest.model_validate(
|
|
jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
|
|
)
|
|
assert isinstance(request.root, InitializeRequest)
|
|
|
|
result = ServerResult(
|
|
InitializeResult(
|
|
protocolVersion=LATEST_PROTOCOL_VERSION,
|
|
capabilities=ServerCapabilities(
|
|
logging=None,
|
|
resources=None,
|
|
tools=None,
|
|
experimental=None,
|
|
prompts=None,
|
|
),
|
|
serverInfo=Implementation(name="mock-server", version="0.1.0"),
|
|
instructions="The server instructions.",
|
|
)
|
|
)
|
|
|
|
async with server_to_client_send:
|
|
assert isinstance(jsonrpc_request.message.root, JSONRPCRequest)
|
|
await server_to_client_send.send(
|
|
MessageFrame(
|
|
message=JSONRPCMessage(
|
|
JSONRPCResponse(
|
|
jsonrpc="2.0",
|
|
id=jsonrpc_request.message.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.message, JSONRPCMessage)
|
|
initialized_notification = ClientNotification.model_validate(
|
|
jsonrpc_notification.message.model_dump(
|
|
by_alias=True, mode="json", exclude_none=True
|
|
)
|
|
)
|
|
|
|
async def listen_session():
|
|
async for message in session.incoming_messages:
|
|
if isinstance(message, Exception):
|
|
raise message
|
|
|
|
async with (
|
|
ClientSession(server_to_client_receive, client_to_server_send) as session,
|
|
anyio.create_task_group() as tg,
|
|
client_to_server_send,
|
|
client_to_server_receive,
|
|
server_to_client_send,
|
|
server_to_client_receive,
|
|
):
|
|
tg.start_soon(mock_server)
|
|
tg.start_soon(listen_session)
|
|
result = await session.initialize()
|
|
|
|
# Assert the result
|
|
assert isinstance(result, InitializeResult)
|
|
assert result.protocolVersion == LATEST_PROTOCOL_VERSION
|
|
assert isinstance(result.capabilities, ServerCapabilities)
|
|
assert result.serverInfo == Implementation(name="mock-server", version="0.1.0")
|
|
assert result.instructions == "The server instructions."
|
|
|
|
# Check that the client sent the initialized notification
|
|
assert initialized_notification
|
|
assert isinstance(initialized_notification.root, InitializedNotification)
|