Make MCPConnectionManager and MCPAggregator close/exit more resilient to distributed environments

Author: saqadri

src/mcp_agent/mcp/mcp_aggregator.py

@@ -231,25 +231,29 @@ async def close(self):
                                         None, None, None
                                     )
                                 except Exception as e:
-                                    logger.error(
-                                        f"Error during connection manager __aexit__: {e}"
+                                    logger.warning(
+                                        f"Error during connection manager cleanup: {e}"
                                     )
 
                                 # Clean up the connection manager from the context
                                 delattr(self.context, "_mcp_connection_manager")
                                 logger.info(
                                     "Connection manager successfully closed and removed from context"
                                 )
-
-                self.initialized = False
+                        else:
+                            logger.debug(
+                                f"Aggregator closing with ref count {current_count}, "
+                                "connection manager will remain active"
+                            )
             except Exception as e:
                 logger.error(
                     f"Error during connection manager cleanup: {e}", exc_info=True
                 )
-                # TODO: saqadri (FA1) - Even if there's an error, we should mark ourselves as uninitialized
-                self.initialized = False
                 span.set_status(trace.Status(trace.StatusCode.ERROR))
                 span.record_exception(e)
+            finally:
+                # Always mark as uninitialized regardless of errors
+                self.initialized = False
 
     @classmethod
     async def create(
@@ -289,7 +293,12 @@ async def create(
                 logger.error(f"Error creating MCPAggregator: {e}")
                 span.set_status(trace.Status(trace.StatusCode.ERROR))
                 span.record_exception(e)
-                await instance.__aexit__(None, None, None)
+                try:
+                    await instance.__aexit__(None, None, None)
+                except Exception as cleanup_error:
+                    logger.warning(
+                        f"Error during MCPAggregator cleanup: {cleanup_error}"
+                    )
 
     async def load_server(self, server_name: str):
         """

src/mcp_agent/mcp/mcp_connection_manager.py

@@ -252,13 +252,19 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
 
             # Then close the task group if it's active
             if self._tg_active:
-                await self._tg.__aexit__(exc_type, exc_val, exc_tb)
-                self._tg_active = False
-                self._tg = None
+                try:
+                    await self._tg.__aexit__(exc_type, exc_val, exc_tb)
+                except Exception as e:
+                    logger.warning(
+                        f"MCPConnectionManager: Error during task group cleanup: {e}"
+                    )
+                finally:
+                    self._tg_active = False
+                    self._tg = None
         except AttributeError:  # Handle missing `_exceptions`
             pass
         except Exception as e:
-            logger.error(f"MCPConnectionManager: Error during shutdown: {e}")
+            logger.warning(f"MCPConnectionManager: Error during shutdown: {e}")
 
     async def launch_server(
         self,