refactor: add abstract profile manager

Author: JoJoJoJoJoJoJo

src/mcpm/core/mcp/client_connection.py

@@ -0,0 +1,69 @@
+import asyncio
+import logging
+from typing import Optional, cast
+
+from mcp import ClientSession, InitializeResult, StdioServerParameters, stdio_client
+from mcp.client.sse import sse_client
+
+from mcpm.core.schema import ServerConfig, SSEServerConfig, STDIOServerConfig
+
+logger = logging.getLogger(__name__)
+
+
+def _stdio_transport_context(server_config: ServerConfig):
+    server_config = cast(STDIOServerConfig, server_config)
+    server_params = StdioServerParameters(command=server_config.command, args=server_config.args, env=server_config.env)
+    return stdio_client(server_params)
+
+
+def _sse_transport_context(server_config: ServerConfig):
+    server_config = cast(SSEServerConfig, server_config)
+    return sse_client(server_config.url, headers=server_config.headers)
+
+
+class ServerConnection:
+    def __init__(self, server_config: ServerConfig) -> None:
+        self.session: Optional[ClientSession] = None
+        self.session_initialized_response: Optional[InitializeResult] = None
+        self._initialized = False
+        self.server_config = server_config
+        self._initialized_event = asyncio.Event()
+        self._shutdown_event = asyncio.Event()
+
+        self._transport_context_factory = (
+            _stdio_transport_context if isinstance(server_config, STDIOServerConfig) else _sse_transport_context
+        )
+
+        self._server_task = asyncio.create_task(self._server_lifespan_cycle())
+
+    def healthy(self) -> bool:
+        return self.session is not None and self._initialized
+
+    # block until client session is initialized
+    async def wait_for_initialization(self):
+        await self._initialized_event.wait()
+
+    # request for client session to gracefully close
+    async def request_for_shutdown(self):
+        self._shutdown_event.set()
+
+    # block until client session is shutdown
+    async def wait_for_shutdown_request(self):
+        await self._shutdown_event.wait()
+
+    async def _server_lifespan_cycle(self):
+        try:
+            async with self._transport_context_factory(self.server_config) as (read, write):
+                async with ClientSession(read, write) as session:
+                    self.session_initialized_response = await session.initialize()
+
+                    self.session = session
+                    self._initialized = True
+                    self._initialized_event.set()
+                    # block here so that the session will not be closed after exit scope
+                    # we could retrieve alive session through self.session
+                    await self.wait_for_shutdown_request()
+        except Exception as e:
+            logger.error(f"Failed to connect to server {self.server_config.name}: {e}")
+            self._initialized_event.set()
+            self._shutdown_event.set()

src/mcpm/core/profile/local_profile_manager.py

@@ -0,0 +1,103 @@
+import json
+import os
+from typing import Dict, Optional, override
+
+from mcpm.core.mcp.client_connection import ServerConnection
+from mcpm.core.profile.profile_manager import AbstractProfileManager
+from mcpm.core.schema import Profile, ServerConfig
+
+DEFAULT_PROFILE_PATH = os.path.expanduser("~/.config/mcpm/profiles.json")
+
+
+class LocalProfileManager(AbstractProfileManager):
+    def __init__(self, profile_path: str = DEFAULT_PROFILE_PATH):
+        super().__init__()
+        self.profile_path = profile_path
+        self._profiles = self._load_profiles()
+        self.server_connections: Dict[str, Dict[str, ServerConnection]] = {}
+
+    def _load_profiles(self) -> Dict[str, Profile]:
+        if not os.path.exists(self.profile_path):
+            return {}
+        with open(self.profile_path, "r") as f:
+            conf = json.load(f)
+            return {name: Profile.model_validate(config) for name, config in conf.items()}
+
+    def _save_profiles(self) -> None:
+        with open(self.profile_path, "w") as f:
+            json.dump(self._profiles, f, indent=2)
+
+    def get_profile(self, name: str) -> Optional[Profile]:
+        return self._profiles.get(name)
+
+    def validate_api_key(self, profile_name: str, api_key: str) -> bool:
+        profile = self.get_profile(profile_name)
+        if not profile:
+            return False
+        return api_key == profile.api_key
+
+    def create_profile(self, profile_name: str) -> None:
+        if profile_name in self._profiles:
+            return
+        self._profiles[profile_name] = Profile(name=profile_name, servers=[])
+        self._save_profiles()
+
+    def delete_profile(self, profile_name: str) -> None:
+        if profile_name not in self._profiles:
+            return
+        del self._profiles[profile_name]
+        self._save_profiles()
+
+    def list_profiles(self) -> list[Profile]:
+        return list(self._profiles.values())
+
+    def rename_profile(self, old_name: str, new_name: str) -> None:
+        if old_name not in self._profiles:
+            return
+        if new_name in self._profiles:
+            return
+        self._profiles[new_name] = self._profiles.pop(old_name)
+        self._save_profiles()
+
+    def add_server(self, profile_name: str, server_config: ServerConfig) -> None:
+        if profile_name not in self._profiles:
+            return
+        self._profiles[profile_name].servers.append(server_config)
+        self._save_profiles()
+
+    def remove_server(self, profile_name: str, server_name: str) -> None:
+        if profile_name not in self._profiles:
+            return
+        self._profiles[profile_name].servers = [
+            server for server in self._profiles[profile_name].servers if server.name != server_name
+        ]
+        self._save_profiles()
+
+    def update_server(self, profile_name: str, server_config: ServerConfig) -> None:
+        if profile_name not in self._profiles:
+            return
+        for idx, server in enumerate(self._profiles[profile_name].servers):
+            if server.name == server_config.name:
+                self._profiles[profile_name].servers[idx] = server_config
+                break
+        self._save_profiles()
+
+    @override
+    async def activate_profile(self, profile_name: str) -> Dict[str, ServerConnection]:
+        if profile_name not in self._profiles:
+            return {}
+        servers = self._profiles[profile_name].servers
+        self.server_connections[profile_name] = {
+            server_config.name: ServerConnection(server_config) for server_config in servers
+        }
+        return self.server_connections[profile_name]
+
+    @override
+    async def deactivate_profile(self, profile_name: str) -> None:
+        if profile_name not in self._profiles:
+            return
+        if profile_name not in self.server_connections:
+            return
+        for _, client in self.server_connections[profile_name].items():
+            await client.request_for_shutdown()
+        self.server_connections.pop(profile_name)

src/mcpm/core/profile/profile_manager.py

@@ -0,0 +1,41 @@
+from profile import Profile
+
+from mcpm.core.schema import ServerConfig
+
+
+class AbstractProfileManager:
+    def __init__(self):
+        pass
+
+    def get_profile(self, name: str) -> Profile | None:
+        raise NotImplementedError
+
+    def validate_api_key(self, profile_name: str, api_key: str) -> bool:
+        raise NotImplementedError
+
+    def create_profile(self, profile_name: str) -> None:
+        raise NotImplementedError
+
+    def delete_profile(self, profile_name: str) -> None:
+        raise NotImplementedError
+
+    def list_profiles(self) -> list[Profile]:
+        raise NotImplementedError
+
+    def rename_profile(self, old_name: str, new_name: str) -> None:
+        raise NotImplementedError
+
+    def add_server(self, profile_name: str, server_config: ServerConfig) -> None:
+        raise NotImplementedError
+
+    def remove_server(self, profile_name: str, server_name: str) -> None:
+        raise NotImplementedError
+
+    def update_server(self, profile_name: str, server_config: ServerConfig) -> None:
+        raise NotImplementedError
+
+    async def activate_profile(self, profile_name: str) -> None:
+        raise NotImplementedError
+
+    async def deactivate_profile(self, profile_name: str) -> None:
+        raise NotImplementedError

src/mcpm/core/schema.py

@@ -80,5 +80,5 @@ def to_mcp_proxy_stdio(self) -> STDIOServerConfig:
 
 class Profile(BaseModel):
     name: str
-    api_key: Optional[str]
+    api_key: Optional[str] = None
     servers: list[ServerConfig]

src/mcpm/schemas/server_config.py

@@ -4,4 +4,6 @@
 
 __all__ = ["ServerConfig", "SSEServerConfig", "STDIOServerConfig"]
 
-warnings.warn("mcpm.schemas.server_config is deprecated, use mcpm.core.schema instead", DeprecationWarning)
+warnings.warn(
+    "mcpm.schemas.server_config is deprecated, use mcpm.core.schema instead", DeprecationWarning, stacklevel=2
+)