Fix MCPClient close cancel-scope & add per-server close methods

tinyagent/mcp_client.py

@@ -105,22 +105,24 @@ async def call_tool(self, name: str, arguments: dict):
             raise
 
     async def close(self):
-        """Clean up subprocess and streams."""
-        if self.exit_stack:
-            try:
-                await self.exit_stack.aclose()
-            except (RuntimeError, asyncio.CancelledError) as e:
-                # Log the error but don't re-raise it
-                self.logger.error(f"Error during client cleanup: {e}")
-            finally:
-                # Always reset these regardless of success or failure
-                self.session = None
-                self.exit_stack = AsyncExitStack()
+        """
+        Close this client’s contexts exactly once,
+        and only if we’re in the same cancel scope.
+        """
+        if self._closed:
+            self.logger.debug("MCPClient.close called twice; ignoring.")
+            return
 
-async def run_example():
-    """Example usage of MCPClient with proper logging."""
-    import sys
-    from tinyagent.hooks.logging_manager import LoggingManager
+        if get_cancel_scope() is not self._own_cancel_scope:
+            self.logger.warning(
+                "Closing outside original cancel scope; results may be undefined.",
+            )
+        try:
+            await self.exit_stack.aclose()
+        except Exception as e:
+            self.logger.error(f"Error during client cleanup: {e}")
+        finally:
+            self._closed = True
 
     # Create and configure logging manager
     log_manager = LoggingManager(default_level=logging.INFO)

tinyagent/tiny_agent.py

@@ -1001,3 +1001,18 @@ async def run_example():
     # Clean up
     await agent.close()
     agent_logger.debug("Example completed")
+
+    async def close_server_by_index(self, index: int):
+        if index < 0 or index >= len(self.mcp_clients):
+            raise IndexError(f"No server at index {index}")
+        client = self.mcp_clients.pop(index)
+        self.tool_to_client = {n: c for n, c in self.tool_to_client.items() if c is not client}
+        await client.close()
+
+    async def close_server_for_tool(self, tool_name: str):
+        client = self.tool_to_client.get(tool_name)
+        if not client:
+            raise KeyError(f"No server for tool {tool_name}")
+        self.mcp_clients.remove(client)
+        self.tool_to_client = {n: c for n, c in self.tool_to_client.items() if c is not client}
+        await client.close()