feat!: Use new MCP SessionMessage for Lambda client transport

Author: clareliguori

src/python/src/mcp_lambda/client/lambda_client.py

@@ -6,6 +6,7 @@
 import mcp.types as types
 from aiobotocore.session import get_session
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from mcp.shared.message import SessionMessage
 from pydantic import BaseModel
 
 
@@ -23,11 +24,11 @@ async def lambda_function_client(lambda_function: LambdaFunctionParameters):
     Client transport for Lambda functions: this will invoke a Lambda function
     when requests are sent to the client.
     """
-    read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
-    read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
+    read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
+    read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
 
-    write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
-    write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
+    write_stream: MemoryObjectSendStream[SessionMessage]
+    write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
 
     read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
     write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
@@ -44,7 +45,8 @@ async def invoke_function():
                 "lambda", region_name=lambda_function.region_name
             ) as lambda_client:
                 async with write_stream_reader:
-                    async for message in write_stream_reader:
+                    async for session_message in write_stream_reader:
+                        message = session_message.message
                         logging.debug(
                             f"MCP JSON RPC message raw: {message.__class__.__name__} {message}"
                         )
@@ -113,12 +115,15 @@ async def invoke_function():
                                         ),
                                     )
                                 )
-                                await read_stream_writer.send(error_message)
+                                await read_stream_writer.send(
+                                    SessionMessage(error_message)
+                                )
                             else:
                                 await read_stream_writer.send(exc)
                             continue
 
-                        await read_stream_writer.send(response_message)
+                        session_message = SessionMessage(response_message)
+                        await read_stream_writer.send(session_message)
         except anyio.ClosedResourceError:
             await anyio.lowlevel.checkpoint()
         except Exception as exc:

src/python/tests/client/test_lambda_client.py

@@ -5,6 +5,7 @@
 
 import anyio
 import pytest
+from mcp.shared.message import SessionMessage
 from mcp.types import (
     JSONRPCError,
     JSONRPCMessage,
@@ -56,8 +57,8 @@ async def test_lambda_function_client_success(mock_client_creator, mock_session)
     )
 
     # Create a test message
-    test_message = JSONRPCMessage(
-        root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
+    test_message = SessionMessage(
+        JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping"))
     )
 
     async with lambda_function_client(lambda_parameters) as (read_stream, write_stream):
@@ -70,10 +71,11 @@ async def test_lambda_function_client_success(mock_client_creator, mock_session)
             response = await read_stream.receive()
 
         # Verify the response
-        assert isinstance(response, JSONRPCMessage)
-        assert isinstance(response.root, JSONRPCResponse)
-        assert response.root.id == "response-id"
-        assert response.root.result == {"message": "success"}
+        assert isinstance(response, SessionMessage)
+        assert isinstance(response.message, JSONRPCMessage)
+        assert isinstance(response.message.root, JSONRPCResponse)
+        assert response.message.root.id == "response-id"
+        assert response.message.root.result == {"message": "success"}
 
         # Verify Lambda was invoked with correct parameters
         mock_client.invoke.assert_called_once()
@@ -100,8 +102,10 @@ async def test_lambda_function_notification_success(mock_client_creator, mock_se
     )
 
     # Create a test message
-    test_message = JSONRPCMessage(
-        root=JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized")
+    test_message = SessionMessage(
+        JSONRPCMessage(
+            root=JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized")
+        )
     )
 
     async with lambda_function_client(lambda_parameters) as (read_stream, write_stream):
@@ -146,9 +150,11 @@ async def test_lambda_function_client_function_error(mock_client_creator, mock_s
     )
 
     # Create a test message
-    test_message = JSONRPCMessage(
-        root=JSONRPCRequest(
-            jsonrpc="2.0", id=1, method="call/tool", params={"hello": "world"}
+    test_message = SessionMessage(
+        JSONRPCMessage(
+            root=JSONRPCRequest(
+                jsonrpc="2.0", id=1, method="call/tool", params={"hello": "world"}
+            )
         )
     )
 
@@ -162,13 +168,15 @@ async def test_lambda_function_client_function_error(mock_client_creator, mock_s
             response = await read_stream.receive()
 
         # Verify the response is an error message
-        assert isinstance(response, JSONRPCMessage)
-        assert isinstance(response.root, JSONRPCError)
-        assert response.root.id == 1
-        assert response.root.error is not None
-        assert response.root.error.code == 500
+        assert isinstance(response, SessionMessage)
+        assert isinstance(response.message, JSONRPCMessage)
+        assert isinstance(response.message.root, JSONRPCError)
+        assert response.message.root.id == 1
+        assert response.message.root.error is not None
+        assert response.message.root.error.code == 500
         assert (
-            "Function invoke returned a function error" in response.root.error.message
+            "Function invoke returned a function error"
+            in response.message.root.error.message
         )
 
         # Verify Lambda was invoked with correct parameters
@@ -198,8 +206,8 @@ async def test_lambda_function_client_invoke_exception(
     mock_client.invoke.side_effect = Exception("Connection error")
 
     # Create a test message
-    test_message = JSONRPCMessage(
-        root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
+    test_message = SessionMessage(
+        JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping"))
     )
 
     async with lambda_function_client(lambda_parameters) as (read_stream, write_stream):
@@ -212,12 +220,13 @@ async def test_lambda_function_client_invoke_exception(
             response = await read_stream.receive()
 
         # Verify the response is an error message
-        assert isinstance(response, JSONRPCMessage)
-        assert isinstance(response.root, JSONRPCError)
-        assert response.root.id == 1
-        assert response.root.error is not None
-        assert response.root.error.code == 500
-        assert "Connection error" in response.root.error.message
+        assert isinstance(response, SessionMessage)
+        assert isinstance(response.message, JSONRPCMessage)
+        assert isinstance(response.message.root, JSONRPCError)
+        assert response.message.root.id == 1
+        assert response.message.root.error is not None
+        assert response.message.root.error.code == 500
+        assert "Connection error" in response.message.root.error.message
 
         # Verify Lambda was invoked with correct parameters
         mock_client.invoke.assert_called_once()
@@ -246,8 +255,8 @@ async def test_lambda_function_client_invalid_response(
     )
 
     # Create a test message
-    test_message = JSONRPCMessage(
-        root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
+    test_message = SessionMessage(
+        JSONRPCMessage(root=JSONRPCRequest(jsonrpc="2.0", id=1, method="ping"))
     )
 
     async with lambda_function_client(lambda_parameters) as (read_stream, write_stream):
@@ -260,12 +269,16 @@ async def test_lambda_function_client_invalid_response(
             response = await read_stream.receive()
 
         # Verify the response is an error message
-        assert isinstance(response, JSONRPCMessage)
-        assert isinstance(response.root, JSONRPCError)
-        assert response.root.id == 1
-        assert response.root.error is not None
-        assert response.root.error.code == 500
-        assert "4 validation errors for JSONRPCMessage" in response.root.error.message
+        assert isinstance(response, SessionMessage)
+        assert isinstance(response.message, JSONRPCMessage)
+        assert isinstance(response.message.root, JSONRPCError)
+        assert response.message.root.id == 1
+        assert response.message.root.error is not None
+        assert response.message.root.error.code == 500
+        assert (
+            "4 validation errors for JSONRPCMessage"
+            in response.message.root.error.message
+        )
 
         # Verify Lambda was invoked with correct parameters
         mock_client.invoke.assert_called_once()