@@ -42,39 +42,32 @@ async def websocket_client(url: str) -> AsyncGenerator[
42
42
43
43
# Connect using websockets, requesting the "mcp" subprotocol
44
44
async with ws_connect (url , subprotocols = [Subprotocol ("mcp" )]) as ws :
45
+
45
46
async def ws_reader ():
46
47
"""
47
48
Reads text messages from the WebSocket, parses them as JSON-RPC messages,
48
49
and sends them into read_stream_writer.
49
50
"""
50
- try :
51
- async with read_stream_writer :
52
- async for raw_text in ws :
53
- try :
54
- message = types .JSONRPCMessage .model_validate_json (raw_text )
55
- await read_stream_writer .send (message )
56
- except ValidationError as exc :
57
- # If JSON parse or model validation fails, send the exception
58
- await read_stream_writer .send (exc )
59
- except (anyio .ClosedResourceError , Exception ) as e :
60
- await ws .close ()
61
- raise e
51
+ async with read_stream_writer :
52
+ async for raw_text in ws :
53
+ try :
54
+ message = types .JSONRPCMessage .model_validate_json (raw_text )
55
+ await read_stream_writer .send (message )
56
+ except ValidationError as exc :
57
+ # If JSON parse or model validation fails, send the exception
58
+ await read_stream_writer .send (exc )
62
59
63
60
async def ws_writer ():
64
61
"""
65
62
Reads JSON-RPC messages from write_stream_reader and sends them to the server.
66
63
"""
67
- try :
68
- async with write_stream_reader :
69
- async for message in write_stream_reader :
70
- # Convert to a dict, then to JSON
71
- msg_dict = message .model_dump (
72
- by_alias = True , mode = "json" , exclude_none = True
73
- )
74
- await ws .send (json .dumps (msg_dict ))
75
- except (anyio .ClosedResourceError , Exception ) as e :
76
- await ws .close ()
77
- raise e
64
+ async with write_stream_reader :
65
+ async for message in write_stream_reader :
66
+ # Convert to a dict, then to JSON
67
+ msg_dict = message .model_dump (
68
+ by_alias = True , mode = "json" , exclude_none = True
69
+ )
70
+ await ws .send (json .dumps (msg_dict ))
78
71
79
72
async with anyio .create_task_group () as tg :
80
73
# Start reader and writer tasks
0 commit comments