@@ -220,33 +220,58 @@ async def connect(self) -> "MCPToolkit":
220
220
self ._exit_stack = AsyncExitStack ()
221
221
222
222
try :
223
- # Connect to all clients using AsyncExitStack
224
- for i , client in enumerate (self .clients ):
225
- try :
226
- # Use MCPClient directly as async context manager
227
- await self ._exit_stack .enter_async_context (client )
228
- msg = f"Connected to client { i + 1 } /{ len (self .clients )} "
229
- logger .debug (msg )
230
- except Exception as e :
231
- logger .error (f"Failed to connect to client { i + 1 } : { e } " )
232
- # AsyncExitStack will handle cleanup of already connected
233
- await self ._exit_stack .aclose ()
234
- self ._exit_stack = None
235
- error_msg = f"Failed to connect to client { i + 1 } : { e } "
236
- raise MCPConnectionError (error_msg ) from e
223
+ # Apply timeout to the entire connection process
224
+ import asyncio
225
+
226
+ timeout_seconds = self .timeout or 30.0
227
+ await asyncio .wait_for (
228
+ self ._connect_all_clients (), timeout = timeout_seconds
229
+ )
237
230
238
231
self ._is_connected = True
239
232
msg = f"Successfully connected to { len (self .clients )} MCP servers"
240
233
logger .info (msg )
241
234
return self
242
235
236
+ except (asyncio .TimeoutError , asyncio .CancelledError ):
237
+ self ._is_connected = False
238
+ if self ._exit_stack :
239
+ await self ._exit_stack .aclose ()
240
+ self ._exit_stack = None
241
+
242
+ timeout_seconds = self .timeout or 30.0
243
+ error_msg = (
244
+ f"Connection timeout after { timeout_seconds } s. "
245
+ f"One or more MCP servers are not responding. "
246
+ f"Please check if the servers are running and accessible."
247
+ )
248
+ logger .error (error_msg )
249
+ raise MCPConnectionError (error_msg )
250
+
243
251
except Exception :
244
252
self ._is_connected = False
245
253
if self ._exit_stack :
246
254
await self ._exit_stack .aclose ()
247
255
self ._exit_stack = None
248
256
raise
249
257
258
+ async def _connect_all_clients (self ):
259
+ r"""Connect to all clients sequentially."""
260
+ # Connect to all clients using AsyncExitStack
261
+ for i , client in enumerate (self .clients ):
262
+ try :
263
+ # Use MCPClient directly as async context manager
264
+ await self ._exit_stack .enter_async_context (client )
265
+ msg = f"Connected to client { i + 1 } /{ len (self .clients )} "
266
+ logger .debug (msg )
267
+ except Exception as e :
268
+ logger .error (f"Failed to connect to client { i + 1 } : { e } " )
269
+ # AsyncExitStack will cleanup already connected clients
270
+ await self ._exit_stack .aclose ()
271
+ self ._exit_stack = None
272
+ error_msg = f"Failed to connect to client { i + 1 } : { e } "
273
+ raise MCPConnectionError (error_msg ) from e
274
+
250
275
async def disconnect (self ):
251
276
r"""Disconnect from all MCP servers."""
252
277
if not self ._is_connected :
0 commit comments