optimize cleanup function.

Author: JDOxygen

docs/docs_en/api/mas.md

@@ -43,7 +43,7 @@ oxygent/mas.py
 | `init_all_oxy()` | Yes | `None` | Initialize all registered Oxy objects |
 | `batch_init_oxy()` | Yes | `None` | Batch initialize oxy objects of specified types |
 | `create_vearch_table()` | Yes | `None` | Create Vearch tables for tools |
-| `cleanup_servers()` | Yes | `None` | Gracefully shut down remote servers/clients |
+| `cleanup_all()` | Yes | `None` | Gracefully release resources held by all registered Oxy components |
 | `add_oxy()` | No | `None` | Register a single Oxy object |
 | `add_oxy_list()` | No | `None` | Register a list of Oxy objects |
 | `call()` | Yes | `Any` | Invoke an Oxy component directly and return its output |

docs/docs_zh/api/mas.md

@@ -43,7 +43,7 @@ oxygent/mas.py
 | `init_all_oxy()` | 是 | `None` | 初始化所有已注册的 Oxy 对象 |
 | `batch_init_oxy()` | 是 | `None` | 批量初始化指定类型的 Oxy 对象 |
 | `create_vearch_table()` | 是 | `None` | 为工具创建 Vearch 表 |
-| `cleanup_servers()` | 是 | `None` | 优雅地关闭远程服务器/客户端 |
+| `cleanup_all()` | 是 | `None` | 优雅地释放所有已注册 Oxy 组件持有的资源 |
 | `add_oxy()` | 否 | `None` | 注册单个 Oxy 对象 |
 | `add_oxy_list()` | 否 | `None` | 注册一组 Oxy 对象 |
 | `call()` | 是 | `Any` | 直接调用一个 Oxy 组件并返回其输出 |

oxygent/databases/db_es/jes_es.py

@@ -124,7 +124,8 @@ async def delete(self, index_name: str, doc_id: str) -> dict[str, Any]:
 
     async def close(self) -> None:
         """Close the Elasticsearch connection."""
-        return await self._run_sync(self.client.close)
+        if self.client is not None:
+            return await self._run_sync(self.client.close)
 
 
 async def main() -> None:

oxygent/embedding_cache.py

@@ -123,8 +123,16 @@ def load(self) -> dict[str, Any]:
         """Load the on-disk cache if it exists; otherwise, return an empty dict."""
         if not os.path.exists(self.file):
             return dict()
-        with open(self.file, "rb") as f:
-            return pickle.load(f)
+        try:
+            with open(self.file, "rb") as f:
+                return pickle.load(f)
+        except Exception as e:
+            logger.warning(
+                f"Failed to load embedding cache from {self.file}, "
+                f"starting with empty cache: {e}",
+                exc_info=True,
+            )
+            return dict()
 
     # TODO: save embeddings
     def save(self) -> None:

oxygent/live_prompt/__init__.py

@@ -4,7 +4,12 @@
 """
 
 # Core ES-based prompt management
-from .manager import PromptManager, get_dynamic_prompt, get_prompt_manager
+from .manager import (
+    PromptManager,
+    close_prompt_manager,
+    get_dynamic_prompt,
+    get_prompt_manager,
+)
 
 # Prompt optimization
 from .optimizer import PromptOptimizer, get_prompt_optimizer
@@ -30,6 +35,7 @@
     # Core ES-based prompt management
     "get_prompt_manager",
     "get_dynamic_prompt",
+    "close_prompt_manager",
     "PromptManager",
     # Hot-reload functionality
     "setup_dynamic_agents",

oxygent/live_prompt/manager.py

@@ -740,7 +740,7 @@ async def close_prompt_manager() -> None:
                 logger.info("Version sync stopped during shutdown")
 
             await prompt_manager.close()
-            logger.info("Prompt manager closed successfully")
+            logger.debug("Prompt manager closed successfully")
         except Exception as e:
             logger.error(f"Error closing prompt manager: {e}", exc_info=True)
         finally:

oxygent/mas.py

@@ -41,7 +41,6 @@
 from .oxy.base_flow import BaseFlow
 from .oxy.base_tool import BaseTool
 from .oxy.llms.base_llm import BaseLLM
-from .oxy.mcp_tools.base_mcp_client import BaseMCPClient
 from .routes import router
 from .schemas import OxyRequest, OxyResponse, SSEMessage, WebResponse
 from .utils.common_utils import (
@@ -213,17 +212,41 @@ async def __aenter__(self) -> "MAS":
         return self
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        async def _safe(coro, label: str) -> None:
+            try:
+                await coro
+                logger.info(f"✓ {label} shutdown completed.")
+            except Exception as e:
+                logger.warning(f"Error during shutdown ({label}): {e}", exc_info=True)
+
+        logger.info("=" * 64)
+        logger.info("🔒 OxyGent MAS Application Shutdown Initiated")
+        logger.info("=" * 64)
+
+        # 1. Wait for background tasks
         all_tasks = [t for tasks in self.background_tasks.values() for t in tasks]
         if all_tasks:
-            await asyncio.gather(*all_tasks)
+            await _safe(
+                asyncio.gather(*all_tasks, return_exceptions=True),
+                "background tasks",
+            )
+
+        # 2. Tear down resources — each step runs regardless of prior failures
+        from .live_prompt import close_prompt_manager
+
+        await _safe(close_prompt_manager(), "prompt manager")
+
+        if self.es_client:
+            await _safe(self.es_client.close(), "ES client")
+
+        if self.redis_client:
+            await _safe(self.redis_client.close(), "Redis client")
+
+        await _safe(self.cleanup_all(), "oxy cleanup")
+
         logger.info("=" * 64)
         logger.info("🪂 OxyGent MAS Application Exit")
         logger.info("=" * 64)
-        if self.es_client:
-            await self.es_client.close()
-        if self.redis_client:
-            await self.redis_client.close()
-        await self.cleanup_servers()
 
     def add_background_task(self, trace_id: str, task: asyncio.Task) -> None:
         """Register a background task under the given trace_id."""
@@ -341,24 +364,21 @@ async def init(self) -> None:
             logger.warning(f"Failed to setup dynamic agents: {e}", exc_info=True)
         logger.info("=" * 64)
 
-    async def cleanup_servers(self) -> None:
-        """Gracefully shut down remote servers/clients.
+    async def cleanup_all(self) -> None:
+        """Gracefully release resources held by all registered Oxy components.
 
-        The method concurrently calls ``cleanup()`` on every
-        :class:`BaseMCPClient` that has been registered.  It is automatically
-        invoked by :func:`__aexit__`.
+        Calls ``cleanup()`` on every registered :class:`Oxy` instance.
+        Each cleanup is invoked individually so that a failure in one
+        component does not prevent others from releasing their resources.
         """
-        cleanup_tasks = []
-        for oxy in self.oxy_name_to_oxy.values():
-            if not isinstance(oxy, BaseMCPClient):
-                continue
-            cleanup_tasks.append(asyncio.create_task(oxy.cleanup()))
-
-        if cleanup_tasks:
+        for oxy_name, oxy in self.oxy_name_to_oxy.items():
             try:
-                await asyncio.gather(*cleanup_tasks, return_exceptions=False)
+                await oxy.cleanup()
             except Exception as e:
-                logger.warning(f"Warning during final cleanup: {e}", exc_info=True)
+                logger.warning(
+                    f"Error during cleanup of oxy '{oxy_name}': {e}",
+                    exc_info=True,
+                )
 
     async def init_db(self) -> None:
         """Es --- (table_name: key)
@@ -635,21 +655,59 @@ async def batch_init_oxy(self, *class_type: type) -> None:
             class_types: List of class types to initialize (e.g., BaseLLM, BaseTool, BaseAgent).
 
         NOTE:
-            Initialize all oxy objects of the specified class types,
+            Initialize all oxy objects of the specified class types.
+            If any init fails, already-initialized oxy objects of the same
+            batch are cleaned up before the exception propagates.
         """
-        tasks = []
-        for oxy_name in list(self.oxy_name_to_oxy.keys()):
-            oxy = self.oxy_name_to_oxy[oxy_name]
-            if not isinstance(oxy, class_type):
-                continue
+        targets = [
+            (name, oxy)
+            for name, oxy in self.oxy_name_to_oxy.items()
+            if isinstance(oxy, class_type)
+        ]
+        for _name, oxy in targets:
             oxy.set_mas(self)
-            task = oxy.init()
-            if Config.get_tool_is_concurrent_init():
-                tasks.append(task)
-            else:
-                await task
-        if tasks:
-            await asyncio.gather(*tasks)
+
+        initialized_names: list[str] = []
+        if Config.get_tool_is_concurrent_init():
+            results = await asyncio.gather(
+                *(oxy.init() for _name, oxy in targets),
+                return_exceptions=True,
+            )
+            first_error = None
+            for (name, _oxy), result in zip(targets, results):
+                if isinstance(result, Exception):
+                    logger.error(
+                        f"Failed to initialize oxy '{name}': {result}", exc_info=result
+                    )
+                    if first_error is None:
+                        first_error = result
+                else:
+                    initialized_names.append(name)
+            if first_error is not None:
+                for name in initialized_names:
+                    try:
+                        await self.oxy_name_to_oxy[name].cleanup()
+                    except Exception as ce:
+                        logger.warning(
+                            f"Error during rollback cleanup of oxy '{name}': {ce}",
+                            exc_info=True,
+                        )
+                raise first_error
+        else:
+            try:
+                for name, oxy in targets:
+                    await oxy.init()
+                    initialized_names.append(name)
+            except Exception:
+                for name in initialized_names:
+                    try:
+                        await self.oxy_name_to_oxy[name].cleanup()
+                    except Exception as ce:
+                        logger.warning(
+                            f"Error during rollback cleanup of oxy '{name}': {ce}",
+                            exc_info=True,
+                        )
+                raise
 
     async def init_all_oxy(self) -> None:
         """Initializing all tools and agents assign values of agent.tools to each

oxygent/oxy/agents/a2a_client_agent.py

@@ -239,6 +239,21 @@ async def init(self) -> None:
             extra={"agent": self.name, "card_url": getattr(self._card, "url", "")},
         )
 
+    async def cleanup(self) -> None:
+        """Close the HTTP client and release connection pool resources."""
+        if self._http_client is not None:
+            try:
+                await self._http_client.aclose()
+            except Exception as e:
+                logger.warning(
+                    f"Error closing HTTP client for A2AClientAgent '{self.name}': {e}",
+                    exc_info=True,
+                )
+            finally:
+                self._http_client = None
+                self._client = None
+                self._card = None
+
     @staticmethod
     def _task_state(task: Task | None) -> str:
         if not task or not task.status:

oxygent/oxy/agents/shell_use_agent.py

@@ -32,15 +32,46 @@ async def init(self) -> None:
 
         client = paramiko.SSHClient()
         client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-        client.connect(**self.auth_info)
+        try:
+            client.connect(**self.auth_info)
+        except Exception as e:
+            logger.error(
+                f"Failed to establish SSH connection for '{self.name}': {e}",
+                exc_info=True,
+            )
+            client.close()
+            raise
+
         ssh_channel = client.invoke_shell()
 
         await asyncio.sleep(1)
         if ssh_channel.recv_ready():
             output = ssh_channel.recv(4096).decode()
             self.mas.global_data["hello_terminal"] = clean_ansi_codes(output)
+        self.mas.global_data["ssh_client"] = client
         self.mas.global_data["ssh_channel"] = ssh_channel
 
+    async def cleanup(self) -> None:
+        """Close the SSH channel and client connection."""
+        channel = self.mas.global_data.pop("ssh_channel", None) if self.mas else None
+        client = self.mas.global_data.pop("ssh_client", None) if self.mas else None
+        try:
+            if channel is not None:
+                channel.close()
+        except Exception as e:
+            logger.warning(
+                f"Error closing SSH channel for '{self.name}': {e}",
+                exc_info=True,
+            )
+        try:
+            if client is not None:
+                client.close()
+        except Exception as e:
+            logger.warning(
+                f"Error closing SSH client for '{self.name}': {e}",
+                exc_info=True,
+            )
+
     def _parse_llm_response(
         self, ori_response: str, oxy_request: OxyRequest = None
     ) -> LLMResponse:

oxygent/oxy/agents/sse_oxy_agent.py

@@ -24,18 +24,30 @@ class SSEOxyGent(RemoteAgent):
     async def init(self) -> None:
         """Initialize the SSE remote agent connection."""
         await super().init()
-
-        async with httpx.AsyncClient() as client:
-            response = await client.get(build_url(self.server_url, "/get_organization"))
-            self.org = response.json()["data"]["organization"]
-
-            if self.desc == "":
+        try:
+            async with httpx.AsyncClient() as client:
                 response = await client.get(
-                    build_url(self.server_url, "/get_description")
+                    build_url(self.server_url, "/get_organization")
                 )
-                if response.json().get("code") == 200:
-                    self.desc = response.json()["data"]["description"]
-        self._set_desc_for_llm()
+                self.org = response.json()["data"]["organization"]
+
+                if self.desc == "":
+                    response = await client.get(
+                        build_url(self.server_url, "/get_description")
+                    )
+                    if response.json().get("code") == 200:
+                        self.desc = response.json()["data"]["description"]
+            self._set_desc_for_llm()
+        except Exception as e:
+            logger.error(
+                f"Failed to initialize SSEOxyGent '{self.name}' "
+                f"(url={self.server_url}): {e}",
+                exc_info=True,
+            )
+            raise RuntimeError(
+                f"Failed to initialize SSEOxyGent '{self.name}' "
+                f"at {self.server_url}: {e}"
+            ) from e
 
     async def _execute(self, oxy_request: OxyRequest) -> OxyResponse:
         """Forward the request to the remote SSE endpoint and stream the response."""