@@ -109,19 +109,25 @@ async def connect(self) -> None:
109109 True
110110 """
111111 loop = asyncio .get_running_loop ()
112-
113- # Set up stdin reader
114- reader = asyncio .StreamReader ()
115- protocol = asyncio .StreamReaderProtocol (reader )
116- await loop .connect_read_pipe (lambda : protocol , sys .stdin )
117- self ._stdin_reader = reader
118-
119- # Set up stdout writer
120- transport , protocol = await loop .connect_write_pipe (asyncio .streams .FlowControlMixin , sys .stdout )
121- self ._stdout_writer = asyncio .StreamWriter (transport , protocol , reader , loop )
122-
123- self ._connected = True
124- logger .info ("stdio transport connected" )
112+ try :
113+ reader = asyncio .StreamReader ()
114+ protocol = asyncio .StreamReaderProtocol (reader )
115+ await loop .connect_read_pipe (lambda : protocol , sys .stdin )
116+ self ._stdin_reader = reader
117+
118+ transport , protocol = await loop .connect_write_pipe (asyncio .streams .FlowControlMixin , sys .stdout )
119+ self ._stdout_writer = asyncio .StreamWriter (transport , protocol , reader , loop )
120+
121+ self ._connected = True
122+ logger .info ("stdio transport connected" )
123+ except (ConnectionResetError , BrokenPipeError ) as e :
124+ logger .error (f"Connection lost during stdio setup: { e } " )
125+ self ._connected = False
126+ raise RuntimeError ("Failed to establish stdio transport connection due to large environment or broken pipe." )
127+ except Exception as e :
128+ logger .error (f"Unexpected error during stdio connect: { e } " )
129+ self ._connected = False
130+ raise
125131
126132 async def disconnect (self ) -> None :
127133 """Clean up stdio streams.
@@ -174,8 +180,15 @@ async def send_message(self, message: Dict[str, Any]) -> None:
174180
175181 try :
176182 data = json .dumps (message )
177- self ._stdout_writer .write (f"{ data } \n " .encode ())
183+ encoded = f"{ data } \n " .encode ()
184+ if len (encoded ) > 10_000_000 : # 10MB safeguard
185+ logger .warning ("Message size exceeds 10MB; may cause pipe reset." )
186+ self ._stdout_writer .write (encoded )
178187 await self ._stdout_writer .drain ()
188+ except (ConnectionResetError , BrokenPipeError ) as e :
189+ logger .error (f"Connection lost while sending message: { e } " )
190+ self ._connected = False
191+ raise RuntimeError ("Connection lost while sending message; possible large environment variable overflow." )
179192 except Exception as e :
180193 logger .error (f"Failed to send message: { e } " )
181194 raise
0 commit comments