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

Author: dsp-ant

src/mcp/client/websocket.py

@@ -1,7 +1,7 @@
 import json
 import logging
 from contextlib import asynccontextmanager
-from typing import AsyncGenerator
+from typing import Any, AsyncGenerator
 
 import anyio
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
@@ -10,6 +10,7 @@
 from websockets.typing import Subprotocol
 
 import mcp.types as types
+from mcp.types import MessageFrame
 
 logger = logging.getLogger(__name__)
 
@@ -19,8 +20,8 @@ async def websocket_client(
     url: str,
 ) -> AsyncGenerator[
     tuple[
-        MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
-        MemoryObjectSendStream[types.JSONRPCMessage],
+        MemoryObjectReceiveStream[MessageFrame[Any] | Exception],
+        MemoryObjectSendStream[MessageFrame[Any]],
     ],
     None,
 ]:
@@ -53,7 +54,11 @@ async def ws_reader():
             async with read_stream_writer:
                 async for raw_text in ws:
                     try:
-                        message = types.JSONRPCMessage.model_validate_json(raw_text)
+                        json_message = types.JSONRPCMessage.model_validate_json(
+                            raw_text
+                        )
+                        # Create MessageFrame with JSON message as root
+                        message = MessageFrame(root=json_message, raw=raw_text)
                         await read_stream_writer.send(message)
                     except ValidationError as exc:
                         # If JSON parse or model validation fails, send the exception
@@ -66,8 +71,8 @@ async def ws_writer():
             """
             async with write_stream_reader:
                 async for message in write_stream_reader:
-                    # Convert to a dict, then to JSON
-                    msg_dict = message.model_dump(
+                    # Extract the JSON-RPC message from MessageFrame and convert to JSON
+                    msg_dict = message.root.model_dump(
                         by_alias=True, mode="json", exclude_none=True
                     )
                     await ws.send(json.dumps(msg_dict))