refactor router config

Author: niechen

.cursor/rules/pytest.mdc

@@ -0,0 +1,6 @@
+---
+description: 
+globs: *.py
+alwaysApply: false
+---
+always run pytest at the end of a major change

src/mcpm/__init__.py

@@ -5,7 +5,9 @@
 # Import version from internal module
 # Import router module
 from . import router
+from .router.router import MCPRouter
+from .router.router_config import RouterConfig
 from .version import __version__
 
 # Define what symbols are exported from this package
-__all__ = ["__version__", "router"]
+__all__ = ["__version__", "router", "MCPRouter", "RouterConfig"]

src/mcpm/router/router.py

@@ -24,9 +24,6 @@
 from mcpm.profile.profile_config import ProfileConfigManager
 from mcpm.schemas.server_config import ServerConfig
 from mcpm.utils.config import (
-    DEFAULT_HOST,
-    DEFAULT_PORT,
-    DEFAULT_SHARE_ADDRESS,
     PROMPT_SPLITOR,
     RESOURCE_SPLITOR,
     RESOURCE_TEMPLATE_SPLITOR,
@@ -35,6 +32,7 @@
 )
 
 from .client_connection import ServerConnection
+from .router_config import RouterConfig
 from .transport import RouterSseTransport
 from .watcher import ConfigWatcher
 
@@ -49,15 +47,16 @@ class MCPRouter:
     Example:
         ```python
         # Initialize with a custom API key
-        router = MCPRouter(api_key="your-api-key")
+        router = MCPRouter(router_config=RouterConfig(api_key="your-api-key"))
 
         # Initialize with custom router configuration
-        router_config = {
-            "host": "localhost",
-            "port": 8080,
-            "share_address": "custom.share.address:8080"
-        }
-        router = MCPRouter(api_key="your-api-key", router_config=router_config)
+        router_config = RouterConfig(
+            host="localhost",
+            port=8080,
+            share_address="custom.share.address:8080",
+            api_key="your-api-key"
+        )
+        router = MCPRouter(router_config=router_config)
 
         # Create a global config from the router's configuration
         router.create_global_config()
@@ -68,18 +67,13 @@ def __init__(
         self,
         reload_server: bool = False,
         profile_path: str | None = None,
-        strict: bool = False,
-        api_key: str | None = None,
-        router_config: dict | None = None,
+        router_config: RouterConfig | None = None,
     ) -> None:
         """
         Initialize the router.
 
         :param reload_server: Whether to reload the server when the config changes
         :param profile_path: Path to the profile file
-        :param strict: Whether to use strict mode for duplicated tool name.
-                       If True, raise error when duplicated tool name is found else auto resolve by adding server name prefix
-        :param api_key: Optional API key to use for authentication.
         :param router_config: Optional router configuration to use instead of the global config
         """
         self.server_sessions: t.Dict[str, ServerConnection] = {}
@@ -94,27 +88,27 @@ def __init__(
         self.watcher: Optional[ConfigWatcher] = None
         if reload_server:
             self.watcher = ConfigWatcher(self.profile_manager.profile_path)
-        self.strict: bool = strict
-        self.api_key = api_key
-        self.router_config = router_config
+        self.router_config = router_config if router_config is not None else RouterConfig()
 
     def create_global_config(self) -> None:
         """
         Create a global configuration from the router's configuration.
         This is useful if you want to initialize the router with a config
         but also want that config to be available globally.
         """
-        if self.api_key is not None:
-            config_manager = ConfigManager()
-            # Save the API key to the global config
-            config_manager.save_share_config(api_key=self.api_key)
-
-            # If router_config is provided, save it to the global config
-            if self.router_config is not None:
-                host = self.router_config.get("host", DEFAULT_HOST)
-                port = self.router_config.get("port", DEFAULT_PORT)
-                share_address = self.router_config.get("share_address", DEFAULT_SHARE_ADDRESS)
-                config_manager.save_router_config(host, port, share_address)
+        # Skip if router_config is None or there's no explicit api_key set
+        if self.router_config is None or self.router_config.api_key is None:
+            return
+
+        config_manager = ConfigManager()
+
+        # Save the API key to the global config
+        config_manager.save_share_config(api_key=self.router_config.api_key)
+
+        # Save router configuration to the global config
+        config_manager.save_router_config(
+            self.router_config.host, self.router_config.port, self.router_config.share_address
+        )
 
     def get_unique_servers(self) -> list[ServerConfig]:
         profiles = self.profile_manager.list_profiles()
@@ -191,7 +185,7 @@ async def add_server(self, server_id: str, server_config: ServerConfig) -> None:
                     # To make sure tool name is unique across all servers
                     tool_name = tool.name
                     if tool_name in self.capabilities_to_server_id["tools"]:
-                        if self.strict:
+                        if self.router_config.strict:
                             raise ValueError(
                                 f"Tool {tool_name} already exists. Please use unique tool names across all servers."
                             )
@@ -210,7 +204,7 @@ async def add_server(self, server_id: str, server_config: ServerConfig) -> None:
                     # To make sure prompt name is unique across all servers
                     prompt_name = prompt.name
                     if prompt_name in self.capabilities_to_server_id["prompts"]:
-                        if self.strict:
+                        if self.router_config.strict:
                             raise ValueError(
                                 f"Prompt {prompt_name} already exists. Please use unique prompt names across all servers."
                             )
@@ -229,7 +223,7 @@ async def add_server(self, server_id: str, server_config: ServerConfig) -> None:
                     # To make sure resource URI is unique across all servers
                     resource_uri = resource.uri
                     if str(resource_uri) in self.capabilities_to_server_id["resources"]:
-                        if self.strict:
+                        if self.router_config.strict:
                             raise ValueError(
                                 f"Resource {resource_uri} already exists. Please use unique resource URIs across all servers."
                             )
@@ -256,7 +250,7 @@ async def add_server(self, server_id: str, server_config: ServerConfig) -> None:
                     # To make sure resource template URI is unique across all servers
                     resource_template_uri_template = resource_template.uriTemplate
                     if resource_template_uri_template in self.capabilities_to_server_id["resource_templates"]:
-                        if self.strict:
+                        if self.router_config.strict:
                             raise ValueError(
                                 f"Resource template {resource_template_uri_template} already exists. Please use unique resource template URIs across all servers."
                             )
@@ -564,7 +558,8 @@ async def get_sse_server_app(
         await self.initialize_router()
 
         # Pass the API key to the RouterSseTransport
-        sse = RouterSseTransport("/messages/", api_key=self.api_key)
+        api_key = None if self.router_config is None else self.router_config.api_key
+        sse = RouterSseTransport("/messages/", api_key=api_key)
 
         async def handle_sse(request: Request) -> None:
             async with sse.connect_sse(

src/mcpm/router/router_config.py

@@ -0,0 +1,17 @@
+from typing import Optional
+
+from pydantic import BaseModel
+
+from mcpm.utils.config import DEFAULT_HOST, DEFAULT_PORT, DEFAULT_SHARE_ADDRESS
+
+
+class RouterConfig(BaseModel):
+    """
+    Router configuration model for MCPRouter
+    """
+
+    host: str = DEFAULT_HOST
+    port: int = DEFAULT_PORT
+    share_address: str = DEFAULT_SHARE_ADDRESS
+    api_key: Optional[str] = None
+    strict: bool = False

src/mcpm/router/transport.py

@@ -245,14 +245,11 @@ def _validate_api_key(self, scope: Scope, api_key: str | None) -> bool:
             return True
 
         # If we have a directly provided API key, verify it matches
-        if self.api_key is not None:
-            # If API key doesn't match, return False
-            if api_key != self.api_key:
-                logger.warning("Unauthorized API key")
-                return False
+        if api_key == self.api_key:
             return True
 
-        # Otherwise, fall back to the original validation logic
+        # At this point, self.api_key is not None but doesn't match the provided api_key
+        # Let's check if this is a share URL that needs special validation
         try:
             config_manager = ConfigManager()
             host = get_key_from_scope(scope, key_name="host") or ""
@@ -264,10 +261,11 @@ def _validate_api_key(self, scope: Scope, api_key: str | None) -> bool:
             share_host_name = urlsplit(share_config["url"]).hostname
             if share_config["url"] and (host_name == share_host_name or host_name != router_config["host"]):
                 share_api_key = share_config["api_key"]
-                if api_key != share_api_key:
-                    logger.warning("Unauthorized API key")
-                    return False
+                if api_key == share_api_key:
+                    return True
         except Exception as e:
             logger.error(f"Failed to validate API key: {e}")
-            return False
-        return True
+
+        # If we reach here, the API key is invalid
+        logger.warning("Unauthorized API key")
+        return False

tests/test_router.py

@@ -10,6 +10,7 @@
 
 from mcpm.router.client_connection import ServerConnection
 from mcpm.router.router import MCPRouter
+from mcpm.router.router_config import RouterConfig
 from mcpm.schemas.server_config import SSEServerConfig
 from mcpm.utils.config import TOOL_SPLITOR
 
@@ -53,41 +54,40 @@ async def test_router_init():
     router = MCPRouter()
     assert router.profile_manager is not None
     assert router.watcher is None
-    assert router.strict is False
-    assert router.api_key is None
-    assert router.router_config is None
+    assert router.router_config is not None
+    assert router.router_config.strict is False
 
     # Test with custom values
-    router_config = {"host": "custom-host", "port": 9000}
+    config = RouterConfig(
+        host="custom-host", port=9000, share_address="custom-share-address", api_key="test-api-key", strict=True
+    )
     router = MCPRouter(
         reload_server=True,
-        strict=True,
-        api_key="test-api-key",
-        router_config=router_config,
+        router_config=config,
     )
 
     assert router.watcher is not None
-    assert router.strict is True
-    assert router.api_key == "test-api-key"
-    assert router.router_config == router_config
+    assert router.router_config == config
+    assert router.router_config.api_key == "test-api-key"
+    assert router.router_config.strict is True
 
 
 def test_create_global_config():
     """Test creating a global config from router config"""
-    router_config = {"host": "custom-host", "port": 9000, "share_address": "custom-share-address"}
+    config = RouterConfig(host="custom-host", port=9000, share_address="custom-share-address", api_key="test-api-key")
 
     with patch("mcpm.router.router.ConfigManager") as mock_config_manager:
         mock_instance = Mock()
         mock_config_manager.return_value = mock_instance
 
-        # Test without API key
-        router = MCPRouter(router_config=router_config)
+        # Test without router_config
+        router = MCPRouter()
         router.create_global_config()
         mock_instance.save_share_config.assert_not_called()
         mock_instance.save_router_config.assert_not_called()
 
-        # Test with API key
-        router = MCPRouter(api_key="test-api-key", router_config=router_config)
+        # Test with router_config
+        router = MCPRouter(router_config=config)
         router.create_global_config()
         mock_instance.save_share_config.assert_called_once_with(api_key="test-api-key")
         mock_instance.save_router_config.assert_called_once_with("custom-host", 9000, "custom-share-address")
@@ -144,7 +144,7 @@ async def test_add_server_unhealthy():
 @pytest.mark.asyncio
 async def test_add_server_duplicate_tool_strict():
     """Test adding a server with duplicate tool name in strict mode"""
-    router = MCPRouter(strict=True)
+    router = MCPRouter(router_config=RouterConfig(strict=True))
 
     # Mock get_active_servers to return all server IDs
     def mock_get_active_servers(_profile):
@@ -207,7 +207,7 @@ def mock_get_active_servers(_profile):
 @pytest.mark.asyncio
 async def test_add_server_duplicate_tool_non_strict():
     """Test adding a server with duplicate tool name in non-strict mode"""
-    router = MCPRouter(strict=False)
+    router = MCPRouter(router_config=RouterConfig(strict=False))
 
     # Mock get_active_servers to return all server IDs
     def mock_get_active_servers(_profile):
@@ -430,7 +430,7 @@ async def test_router_sse_transport_with_api_key():
 @pytest.mark.asyncio
 async def test_get_sse_server_app_with_api_key():
     """Test that the API key is passed to RouterSseTransport when creating the server app"""
-    router = MCPRouter(api_key="test-api-key")
+    router = MCPRouter(router_config=RouterConfig(api_key="test-api-key"))
 
     # Patch the RouterSseTransport constructor and get_active_servers method
     with (
@@ -455,7 +455,7 @@ def mock_get_active_servers(_profile):
 @pytest.mark.asyncio
 async def test_get_sse_server_app_without_api_key():
     """Test that None is passed to RouterSseTransport when no API key is provided"""
-    router = MCPRouter()  # No API key
+    router = MCPRouter()  # No API key or router_config
 
     # Patch the RouterSseTransport constructor and get_active_servers method
     with (