Fix bug in stash and pop

niechen · niechen · commit b8d9b84dbbd1 · 2025-03-28T22:09:48.000+08:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -16,6 +16,7 @@ dependencies = [
     "click>=8.1.3",
     "rich>=12.0.0",
     "requests>=2.28.0",
+    "pydantic>=2.5.1",
 ]
 
 [project.scripts]
diff --git a/src/mcpm/clients/claude_desktop.py b/src/mcpm/clients/claude_desktop.py
@@ -136,16 +136,24 @@ def from_claude_desktop_format(cls, server_name: str, client_config: Dict[str, A
                 
         return ServerConfig.from_dict(server_data)
     
-    def _convert_from_client_format(self, server_name: str, client_config: Dict[str, Any]) -> ServerConfig:
+    def _convert_from_client_format(self, server_name: str, client_config: Any) -> ServerConfig:
         """Convert Claude Desktop format to ServerConfig
         
         Args:
             server_name: Name of the server
-            client_config: Claude Desktop-specific configuration
+            client_config: Claude Desktop-specific configuration or ServerConfig object
             
         Returns:
             ServerConfig object
         """
+        # If client_config is already a ServerConfig object, just return it
+        if isinstance(client_config, ServerConfig):
+            # Ensure the name is set correctly
+            if client_config.name != server_name:
+                client_config.name = server_name
+            return client_config
+            
+        # Otherwise, convert from dict format
         return self.from_claude_desktop_format(server_name, client_config)
     
     def disable_server(self, server_name: str) -> bool:
diff --git a/src/mcpm/commands/pop.py b/src/mcpm/commands/pop.py
@@ -1,13 +1,15 @@
-"""
-Pop command for MCPM - restores previously stashed server configuration
+"""Pop command for MCPM - restores previously stashed server configuration
 """
 
+import logging
 import click
 from rich.console import Console
 
 from mcpm.utils.client_registry import ClientRegistry
+from mcpm.utils.config import ConfigManager
 
 console = Console()
+logger = logging.getLogger(__name__)
 
 @click.command()
 @click.argument("server_name")
@@ -32,16 +34,29 @@ def pop(server_name):
         console.print("Please switch to a supported client using 'mcpm client <client-name>'")
         return
     
-    # Check if the server exists in the stashed configurations
-    if not client_manager.is_server_disabled(server_name):
-        console.print(f"[bold red]Error:[/] Server '{server_name}' not found in stashed configurations.")
+    # Access the global config manager
+    config_manager = ConfigManager()
+    
+    # Check if the server is stashed for this client
+    if not config_manager.is_server_stashed(client, server_name):
+        console.print(f"[bold red]Error:[/] Server '{server_name}' not found in stashed configurations for {client_name}.")
+        return
+    
+    # Get the server configuration from global stashed servers
+    server_data = config_manager.pop_server(client, server_name)
+    if not server_data:
+        console.print(f"[bold red]Error:[/] Failed to retrieve stashed configuration for server '{server_name}'.")
         return
     
-    # Pop (re-enable) the server
-    success = client_manager.enable_server(server_name)
+    # Convert the server configuration to the client's format and add it back
+    # to the active servers
+    server_config = client_manager._convert_from_client_format(server_name, server_data)
+    success = client_manager.add_server(server_config)
     
     if success:
         console.print(f"[bold green]Restored[/] MCP server '{server_name}' for {client_name}")
         console.print("Remember to restart the client for changes to take effect.")
     else:
+        # If adding failed, re-stash the server to avoid data loss
+        config_manager.stash_server(client, server_name, server_data)
         console.print(f"[bold red]Failed to restore[/] '{server_name}' for {client_name}.")
diff --git a/src/mcpm/commands/stash.py b/src/mcpm/commands/stash.py
@@ -1,13 +1,15 @@
-"""
-Stash command for MCPM - temporarily stores server configuration aside
+"""Stash command for MCPM - temporarily stores server configuration aside
 """
 
+import logging
 import click
 from rich.console import Console
 
 from mcpm.utils.client_registry import ClientRegistry
+from mcpm.utils.config import ConfigManager
 
 console = Console()
+logger = logging.getLogger(__name__)
 
 @click.command()
 @click.argument("server_name")
@@ -33,19 +35,38 @@ def stash(server_name):
         return
     
     # Check if the server exists in the active client
-    server_info = client_manager.get_server(server_name)
-    if not server_info:
+    server_config = client_manager.get_server_config(server_name)
+    if not server_config:
         console.print(f"[bold red]Error:[/] Server '{server_name}' not found in {client_name}.")
         return
+        
+    # Convert ServerConfig to dictionary for storage
+    server_data = server_config.to_dict() if hasattr(server_config, 'to_dict') else server_config
+    
+    # Access the global config manager
+    config_manager = ConfigManager()
+    
+    # Check if server is already stashed
+    if config_manager.is_server_stashed(client, server_name):
+        console.print(f"[bold red]Error:[/] Server '{server_name}' is already stashed for {client_name}.")
+        return
     
     # Display server information before stashing
     console.print(f"\n[bold cyan]Stashing server:[/] {server_name}")
     
-    # Stash the server (disable it)
-    success = client_manager.disable_server(server_name)
+    # Store server in global config's stashed_servers
+    stash_success = config_manager.stash_server(client, server_name, server_data)
     
-    if success:
-        console.print(f"[bold yellow]Stashed[/] MCP server '{server_name}' for {client_name}")
-        console.print("Server configuration is stored and can be restored with 'mcpm pop'.")
+    # Remove the server from the client's configuration
+    if stash_success:
+        remove_success = client_manager.remove_server(server_name)
+        
+        if remove_success:
+            console.print(f"[bold yellow]Stashed[/] MCP server '{server_name}' for {client_name}")
+            console.print("Server configuration is stored and can be restored with 'mcpm pop'.")
+        else:
+            # If removing failed, also remove from stashed servers to avoid issues
+            config_manager.pop_server(client, server_name)
+            console.print(f"[bold red]Failed to remove server from {client_name}. Stashing aborted.")
     else:
         console.print(f"[bold red]Failed to stash[/] '{server_name}' for {client_name}.")
diff --git a/src/mcpm/utils/config.py b/src/mcpm/utils/config.py
@@ -108,3 +108,122 @@ def get_supported_clients(self) -> List[str]:
         # We'll import here to avoid circular imports
         from mcpm.utils.client_registry import ClientRegistry
         return ClientRegistry.get_supported_clients()
+        
+    def stash_server(self, client_name: str, server_name: str, server_config: Any) -> bool:
+        """Store a disabled server configuration in the global config
+        
+        Args:
+            client_name: Name of the client the server belongs to
+            server_name: Name of the server to stash
+            server_config: Server configuration to stash (ServerConfig object or dict)
+            
+        Returns:
+            bool: Success or failure
+        """
+        # Ensure stashed_servers section exists
+        if "stashed_servers" not in self._config:
+            self._config["stashed_servers"] = {}
+            
+        # Ensure client section exists
+        if client_name not in self._config["stashed_servers"]:
+            self._config["stashed_servers"][client_name] = {}
+        
+        # Convert ServerConfig to dict if needed
+        try:
+            # If it's a ServerConfig object with to_dict method
+            if hasattr(server_config, 'to_dict') and callable(server_config.to_dict):
+                server_dict = server_config.to_dict()
+            else:
+                # Assume it's already a dict or JSON serializable
+                server_dict = server_config
+                
+            # Add the server configuration
+            self._config["stashed_servers"][client_name][server_name] = server_dict
+            
+            # Save the config
+            self._save_config()
+            return True
+        except Exception as e:
+            logger.error(f"Failed to save stashed server: {e}")
+            return False
+    
+    def pop_server(self, client_name: str, server_name: str) -> Optional[Dict[str, Any]]:
+        """Retrieve a stashed server configuration from the global config
+        
+        Args:
+            client_name: Name of the client the server belongs to
+            server_name: Name of the server to retrieve
+            
+        Returns:
+            Dict: Server configuration or None if not found
+        """
+        # Check if stashed_servers section exists
+        if "stashed_servers" not in self._config:
+            return None
+            
+        # Check if client section exists
+        if client_name not in self._config["stashed_servers"]:
+            return None
+            
+        # Check if server exists
+        if server_name not in self._config["stashed_servers"][client_name]:
+            return None
+            
+        # Get the server configuration
+        server_config = self._config["stashed_servers"][client_name][server_name]
+        
+        # Remove the server from stashed servers
+        del self._config["stashed_servers"][client_name][server_name]
+        
+        # Clean up empty client section if needed
+        if not self._config["stashed_servers"][client_name]:
+            del self._config["stashed_servers"][client_name]
+            
+        # Clean up empty stashed_servers section if needed
+        if not self._config["stashed_servers"]:
+            del self._config["stashed_servers"]
+            
+        # Save the config
+        self._save_config()
+        
+        return server_config
+    
+    def is_server_stashed(self, client_name: str, server_name: str) -> bool:
+        """Check if a server is stashed in the global config
+        
+        Args:
+            client_name: Name of the client the server belongs to
+            server_name: Name of the server to check
+            
+        Returns:
+            bool: True if server is stashed, False otherwise
+        """
+        # Check if stashed_servers section exists
+        if "stashed_servers" not in self._config:
+            return False
+            
+        # Check if client section exists
+        if client_name not in self._config["stashed_servers"]:
+            return False
+            
+        # Check if server exists
+        return server_name in self._config["stashed_servers"][client_name]
+    
+    def get_stashed_servers(self, client_name: str) -> Dict[str, Dict[str, Any]]:
+        """Get all stashed servers for a client
+        
+        Args:
+            client_name: Name of the client to get stashed servers for
+            
+        Returns:
+            Dict: Dictionary of server configurations by name
+        """
+        # Check if stashed_servers section exists
+        if "stashed_servers" not in self._config:
+            return {}
+            
+        # Check if client section exists
+        if client_name not in self._config["stashed_servers"]:
+            return {}
+            
+        return self._config["stashed_servers"][client_name]
diff --git a/src/mcpm/utils/server_config.py b/src/mcpm/utils/server_config.py

Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,7 @@ dependencies = [`
`16`	`16`	`"click>=8.1.3",`
`17`	`17`	`"rich>=12.0.0",`
`18`	`18`	`"requests>=2.28.0",`
	`19`	`+ "pydantic>=2.5.1",`
`19`	`20`	`]`
`20`	`21`
`21`	`22`	`[project.scripts]`