Wrap JSONRPC messages with SessionMessage for metadata support (#590)

tests/client/test_session.py

@@ -3,6 +3,7 @@ import pytest
 
 import mcp.types as types
 from mcp.client.session import DEFAULT_CLIENT_INFO, ClientSession
+from mcp.shared.message import SessionMessage
 from mcp.shared.session import RequestResponder
 from mcp.types import (
     LATEST_PROTOCOL_VERSION,
@@ -24,10 +25,10 @@ from mcp.types import (
 @pytest.mark.anyio
 async def test_client_session_initialize():
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage
+        SessionMessage
     ](1)
     server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage
+        SessionMessage
     ](1)
 
     initialized_notification = None
@@ -35,7 +36,8 @@ async def test_client_session_initialize():
     async def mock_server():
         nonlocal initialized_notification
 
-        jsonrpc_request = await client_to_server_receive.receive()
+        session_message = await client_to_server_receive.receive()
+        jsonrpc_request = session_message.message
         assert isinstance(jsonrpc_request.root, JSONRPCRequest)
         request = ClientRequest.model_validate(
             jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
@@ -59,17 +61,20 @@ async def test_client_session_initialize():
 
         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
-                        ),
+                SessionMessage(
+                    JSONRPCMessage(
+                        JSONRPCResponse(
+                            jsonrpc="2.0",
+                            id=jsonrpc_request.root.id,
+                            result=result.model_dump(
+                                by_alias=True, mode="json", exclude_none=True
+                            ),
+                        )
                     )
                 )
             )
-            jsonrpc_notification = await client_to_server_receive.receive()
+            session_notification = await client_to_server_receive.receive()
+            jsonrpc_notification = session_notification.message
             assert isinstance(jsonrpc_notification.root, JSONRPCNotification)
             initialized_notification = ClientNotification.model_validate(
                 jsonrpc_notification.model_dump(
@@ -116,10 +121,10 @@ async def test_client_session_initialize():
 @pytest.mark.anyio
 async def test_client_session_custom_client_info():
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage
+        SessionMessage
     ](1)
     server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage
+        SessionMessage
     ](1)
 
     custom_client_info = Implementation(name="test-client", version="1.2.3")
@@ -128,7 +133,8 @@ async def test_client_session_custom_client_info():
     async def mock_server():
         nonlocal received_client_info
 
-        jsonrpc_request = await client_to_server_receive.receive()
+        session_message = await client_to_server_receive.receive()
+        jsonrpc_request = session_message.message
         assert isinstance(jsonrpc_request.root, JSONRPCRequest)
         request = ClientRequest.model_validate(
             jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
@@ -146,13 +152,15 @@ async def test_client_session_custom_client_info():
 
         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
-                        ),
+                SessionMessage(
+                    JSONRPCMessage(
+                        JSONRPCResponse(
+                            jsonrpc="2.0",
+                            id=jsonrpc_request.root.id,
+                            result=result.model_dump(
+                                by_alias=True, mode="json", exclude_none=True
+                            ),
+                        )
                     )
                 )
             )
@@ -181,10 +189,10 @@ async def test_client_session_custom_client_info():
 @pytest.mark.anyio
 async def test_client_session_default_client_info():
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage
+        SessionMessage
     ](1)
     server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage
+        SessionMessage
     ](1)
 
     received_client_info = None
@@ -192,7 +200,8 @@ async def test_client_session_default_client_info():
     async def mock_server():
         nonlocal received_client_info
 
-        jsonrpc_request = await client_to_server_receive.receive()
+        session_message = await client_to_server_receive.receive()
+        jsonrpc_request = session_message.message
         assert isinstance(jsonrpc_request.root, JSONRPCRequest)
         request = ClientRequest.model_validate(
             jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
@@ -210,13 +219,15 @@ async def test_client_session_default_client_info():
 
         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
-                        ),
+                SessionMessage(
+                    JSONRPCMessage(
+                        JSONRPCResponse(
+                            jsonrpc="2.0",
+                            id=jsonrpc_request.root.id,
+                            result=result.model_dump(
+                                by_alias=True, mode="json", exclude_none=True
+                            ),
+                        )
                     )
                 )
             )

tests/client/test_stdio.py

@@ -3,6 +3,7 @@ import shutil
 import pytest
 
 from mcp.client.stdio import StdioServerParameters, stdio_client
+from mcp.shared.message import SessionMessage
 from mcp.types import JSONRPCMessage, JSONRPCRequest, JSONRPCResponse
 
 tee: str = shutil.which("tee")  # type: ignore
@@ -22,7 +23,8 @@ async def test_stdio_client():
 
         async with write_stream:
             for message in messages:
-                await write_stream.send(message)
+                session_message = SessionMessage(message)
+                await write_stream.send(session_message)
 
         read_messages = []
         async with read_stream:
@@ -30,7 +32,7 @@ async def test_stdio_client():
                 if isinstance(message, Exception):
                     raise message
 
-                read_messages.append(message)
+                read_messages.append(message.message)
                 if len(read_messages) == 2:
                     break
 

tests/issues/test_192_request_id.py

@@ -3,6 +3,7 @@ import pytest
 
 from mcp.server.lowlevel import NotificationOptions, Server
 from mcp.server.models import InitializationOptions
+from mcp.shared.message import SessionMessage
 from mcp.types import (
     LATEST_PROTOCOL_VERSION,
     ClientCapabilities,
@@ -64,8 +65,10 @@ async def test_request_id_match() -> None:
             jsonrpc="2.0",
         )
 
-        await client_writer.send(JSONRPCMessage(root=init_req))
-        await server_reader.receive()  # Get init response but don't need to check it
+        await client_writer.send(SessionMessage(JSONRPCMessage(root=init_req)))
+        response = (
+            await server_reader.receive()
+        )  # Get init response but don't need to check it
 
         # Send initialized notification
         initialized_notification = JSONRPCNotification(
@@ -73,21 +76,23 @@ async def test_request_id_match() -> None:
             params=NotificationParams().model_dump(by_alias=True, exclude_none=True),
             jsonrpc="2.0",
         )
-        await client_writer.send(JSONRPCMessage(root=initialized_notification))
+        await client_writer.send(
+            SessionMessage(JSONRPCMessage(root=initialized_notification))
+        )
 
         # Send ping request with custom ID
         ping_request = JSONRPCRequest(
             id=custom_request_id, method="ping", params={}, jsonrpc="2.0"
         )
 
-        await client_writer.send(JSONRPCMessage(root=ping_request))
+        await client_writer.send(SessionMessage(JSONRPCMessage(root=ping_request)))
 
         # Read response
         response = await server_reader.receive()
 
         # Verify response ID matches request ID
         assert (
-            response.root.id == custom_request_id
+            response.message.root.id == custom_request_id
         ), "Response ID should match request ID"
 
         # Cancel server task

tests/server/test_lifespan.py

@@ -10,6 +10,7 @@ from pydantic import TypeAdapter
 from mcp.server.fastmcp import Context, FastMCP
 from mcp.server.lowlevel.server import NotificationOptions, Server
 from mcp.server.models import InitializationOptions
+from mcp.shared.message import SessionMessage
 from mcp.types import (
     ClientCapabilities,
     Implementation,
@@ -82,41 +83,49 @@ async def test_lowlevel_server_lifespan():
             clientInfo=Implementation(name="test-client", version="0.1.0"),
         )
         await send_stream1.send(
-            JSONRPCMessage(
-                root=JSONRPCRequest(
-                    jsonrpc="2.0",
-                    id=1,
-                    method="initialize",
-                    params=TypeAdapter(InitializeRequestParams).dump_python(params),
+            SessionMessage(
+                JSONRPCMessage(
+                    root=JSONRPCRequest(
+                        jsonrpc="2.0",
+                        id=1,
+                        method="initialize",
+                        params=TypeAdapter(InitializeRequestParams).dump_python(params),
+                    )
                 )
             )
         )
         response = await receive_stream2.receive()
+        response = response.message
 
         # Send initialized notification
         await send_stream1.send(
-            JSONRPCMessage(
-                root=JSONRPCNotification(
-                    jsonrpc="2.0",
-                    method="notifications/initialized",
+            SessionMessage(
+                JSONRPCMessage(
+                    root=JSONRPCNotification(
+                        jsonrpc="2.0",
+                        method="notifications/initialized",
+                    )
                 )
             )
         )
 
         # Call the tool to verify lifespan context
         await send_stream1.send(
-            JSONRPCMessage(
-                root=JSONRPCRequest(
-                    jsonrpc="2.0",
-                    id=2,
-                    method="tools/call",
-                    params={"name": "check_lifespan", "arguments": {}},
+            SessionMessage(
+                JSONRPCMessage(
+                    root=JSONRPCRequest(
+                        jsonrpc="2.0",
+                        id=2,
+                        method="tools/call",
+                        params={"name": "check_lifespan", "arguments": {}},
+                    )
                 )
             )
         )
 
         # Get response and verify
         response = await receive_stream2.receive()
+        response = response.message
         assert response.root.result["content"][0]["text"] == "true"
 
         # Cancel server task
@@ -178,41 +187,49 @@ async def test_fastmcp_server_lifespan():
             clientInfo=Implementation(name="test-client", version="0.1.0"),
         )
         await send_stream1.send(
-            JSONRPCMessage(
-                root=JSONRPCRequest(
-                    jsonrpc="2.0",
-                    id=1,
-                    method="initialize",
-                    params=TypeAdapter(InitializeRequestParams).dump_python(params),
+            SessionMessage(
+                JSONRPCMessage(
+                    root=JSONRPCRequest(
+                        jsonrpc="2.0",
+                        id=1,
+                        method="initialize",
+                        params=TypeAdapter(InitializeRequestParams).dump_python(params),
+                    )
                 )
             )
         )
         response = await receive_stream2.receive()
+        response = response.message
 
         # Send initialized notification
         await send_stream1.send(
-            JSONRPCMessage(
-                root=JSONRPCNotification(
-                    jsonrpc="2.0",
-                    method="notifications/initialized",
+            SessionMessage(
+                JSONRPCMessage(
+                    root=JSONRPCNotification(
+                        jsonrpc="2.0",
+                        method="notifications/initialized",
+                    )
                 )
             )
         )
 
         # Call the tool to verify lifespan context
         await send_stream1.send(
-            JSONRPCMessage(
-                root=JSONRPCRequest(
-                    jsonrpc="2.0",
-                    id=2,
-                    method="tools/call",
-                    params={"name": "check_lifespan", "arguments": {}},
+            SessionMessage(
+                JSONRPCMessage(
+                    root=JSONRPCRequest(
+                        jsonrpc="2.0",
+                        id=2,
+                        method="tools/call",
+                        params={"name": "check_lifespan", "arguments": {}},
+                    )
                 )
             )
         )
 
         # Get response and verify
         response = await receive_stream2.receive()
+        response = response.message
         assert response.root.result["content"][0]["text"] == "true"
 
         # Cancel server task

tests/server/test_lowlevel_tool_annotations.py

@@ -8,10 +8,10 @@ from mcp.server import Server
 from mcp.server.lowlevel import NotificationOptions
 from mcp.server.models import InitializationOptions
 from mcp.server.session import ServerSession
+from mcp.shared.message import SessionMessage
 from mcp.shared.session import RequestResponder
 from mcp.types import (
     ClientResult,
-    JSONRPCMessage,
     ServerNotification,
     ServerRequest,
     Tool,
@@ -46,10 +46,10 @@ async def test_lowlevel_server_tool_annotations():
         ]
 
     server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage
+        SessionMessage
     ](10)
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage
+        SessionMessage
     ](10)
 
     # Message handler for client

tests/server/test_session.py

@@ -7,11 +7,11 @@ from mcp.server import Server
 from mcp.server.lowlevel import NotificationOptions
 from mcp.server.models import InitializationOptions
 from mcp.server.session import ServerSession
+from mcp.shared.message import SessionMessage
 from mcp.shared.session import RequestResponder
 from mcp.types import (
     ClientNotification,
     InitializedNotification,
-    JSONRPCMessage,
     PromptsCapability,
     ResourcesCapability,
     ServerCapabilities,
@@ -21,10 +21,10 @@ from mcp.types import (
 @pytest.mark.anyio
 async def test_server_session_initialize():
     server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage
+        SessionMessage
     ](1)
     client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[
-        JSONRPCMessage
+        SessionMessage
     ](1)
 
     # Create a message handler to catch exceptions

tests/server/test_stdio.py

@@ -4,6 +4,7 @@ import anyio
 import pytest
 
 from mcp.server.stdio import stdio_server
+from mcp.shared.message import SessionMessage
 from mcp.types import JSONRPCMessage, JSONRPCRequest, JSONRPCResponse
 
 
@@ -29,7 +30,7 @@ async def test_stdio_server():
             async for message in read_stream:
                 if isinstance(message, Exception):
                     raise message
-                received_messages.append(message)
+                received_messages.append(message.message)
                 if len(received_messages) == 2:
                     break
 
@@ -50,7 +51,8 @@ async def test_stdio_server():
 
         async with write_stream:
             for response in responses:
-                await write_stream.send(response)
+                session_message = SessionMessage(response)
+                await write_stream.send(session_message)
 
     stdout.seek(0)
     output_lines = stdout.readlines()