@@ -42,39 +42,32 @@ async def websocket_client(url: str) -> AsyncGenerator[
4242
4343 # Connect using websockets, requesting the "mcp" subprotocol
4444 async with ws_connect (url , subprotocols = [Subprotocol ("mcp" )]) as ws :
45+
4546 async def ws_reader ():
4647 """
4748 Reads text messages from the WebSocket, parses them as JSON-RPC messages,
4849 and sends them into read_stream_writer.
4950 """
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 )
6259
6360 async def ws_writer ():
6461 """
6562 Reads JSON-RPC messages from write_stream_reader and sends them to the server.
6663 """
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 ))
7871
7972 async with anyio .create_task_group () as tg :
8073 # Start reader and writer tasks
0 commit comments