Add support for custom API key and router configuration in MCPRouter

Author: openhands-agent

README.md

@@ -162,6 +162,30 @@ mcpm router share                 # Share the router to public
 mcpm router unshare               # Unshare the router
 ```
 
+#### Programmatic Usage of MCPRouter
+
+You can also use the `MCPRouter` class programmatically in your Python applications. This is especially useful if you want to integrate MCPM into your own applications or scripts without relying on the global configuration.
+
+```python
+from mcpm.router.router import MCPRouter
+
+# Initialize with a custom API key
+router = MCPRouter(api_key="your-custom-api-key")
+
+# Initialize with custom router configuration
+router_config = {
+    "host": "localhost",
+    "port": 8080,
+    "share_address": "custom.share.address:8080"
+}
+router = MCPRouter(api_key="your-custom-api-key", router_config=router_config)
+
+# Optionally, create a global config from the router's configuration
+router.create_global_config()
+```
+
+See the [examples directory](examples/) for more detailed examples of programmatic usage.
+
 ### 🛠️ Utilities (`util`)
 
 ```bash

examples/custom_api_key_example.py

@@ -0,0 +1,28 @@
+#!/usr/bin/env python
+"""
+Example script demonstrating how to use MCPRouter with a custom API key.
+"""
+
+import asyncio
+from mcpm.router.router import MCPRouter
+
+async def main():
+    # Initialize the router with a custom API key
+    router = MCPRouter(api_key="your-custom-api-key")
+    
+    # Optionally, create a global config from the router's configuration
+    # This will save the API key to the global config file
+    # router.create_global_config()
+    
+    # Initialize the router and start the server
+    app = await router.get_sse_server_app(allow_origins=["*"])
+    
+    # Print a message to indicate that the router is ready
+    print("Router initialized with custom API key")
+    print("You can now use this router without setting up a global config")
+    
+    # In a real application, you would start the server here:
+    # await router.start_sse_server(host="localhost", port=8080, allow_origins=["*"])
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/custom_router_config_example.py

@@ -0,0 +1,46 @@
+#!/usr/bin/env python
+"""
+Example script demonstrating how to use MCPRouter with custom configuration.
+"""
+
+import asyncio
+from mcpm.router.router import MCPRouter
+
+async def main():
+    # Define custom router configuration
+    router_config = {
+        "host": "localhost",
+        "port": 8080,
+        "share_address": "custom.share.address:8080"
+    }
+    
+    # Initialize the router with a custom API key and router configuration
+    router = MCPRouter(
+        api_key="your-custom-api-key",
+        router_config=router_config,
+        # You can also specify other parameters:
+        # reload_server=True,  # Reload the server when the config changes
+        # profile_path="/custom/path/to/profile.json",  # Custom profile path
+        # strict=True,  # Use strict mode for duplicated tool names
+    )
+    
+    # Optionally, create a global config from the router's configuration
+    # This will save both the API key and router configuration to the global config file
+    # router.create_global_config()
+    
+    # Initialize the router and start the server
+    app = await router.get_sse_server_app(allow_origins=["*"])
+    
+    # Print a message to indicate that the router is ready
+    print("Router initialized with custom configuration")
+    print("You can now use this router without setting up a global config")
+    
+    # In a real application, you would start the server here:
+    # await router.start_sse_server(
+    #     host=router_config["host"],
+    #     port=router_config["port"],
+    #     allow_origins=["*"]
+    # )
+
+if __name__ == "__main__":
+    asyncio.run(main())

src/mcpm/router/router.py

@@ -23,7 +23,10 @@
 from mcpm.monitor.event import trace_event
 from mcpm.profile.profile_config import ProfileConfigManager
 from mcpm.schemas.server_config import ServerConfig
-from mcpm.utils.config import PROMPT_SPLITOR, RESOURCE_SPLITOR, RESOURCE_TEMPLATE_SPLITOR, TOOL_SPLITOR
+from mcpm.utils.config import (
+    PROMPT_SPLITOR, RESOURCE_SPLITOR, RESOURCE_TEMPLATE_SPLITOR, TOOL_SPLITOR,
+    ConfigManager, DEFAULT_HOST, DEFAULT_PORT, DEFAULT_SHARE_ADDRESS
+)
 
 from .client_connection import ServerConnection
 from .transport import RouterSseTransport
@@ -36,16 +39,42 @@ class MCPRouter:
     """
     A router that aggregates multiple MCP servers (SSE/STDIO) and
     exposes them as a single SSE server.
+    
+    Example:
+        ```python
+        # Initialize with a custom API key
+        router = MCPRouter(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)
+        
+        # Create a global config from the router's configuration
+        router.create_global_config()
+        ```
     """
 
-    def __init__(self, reload_server: bool = False, profile_path: str | None = None, strict: bool = False) -> None:
+    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
+    ) -> 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] = {}
         self.capabilities_mapping: t.Dict[str, t.Dict[str, t.Any]] = defaultdict(dict)
@@ -60,6 +89,26 @@ def __init__(self, reload_server: bool = False, profile_path: str | None = 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
+        
+    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)
 
     def get_unique_servers(self) -> list[ServerConfig]:
         profiles = self.profile_manager.list_profiles()
@@ -496,7 +545,8 @@ async def get_sse_server_app(
         """
         await self.initialize_router()
 
-        sse = RouterSseTransport("/messages/")
+        # Pass the API key to the RouterSseTransport
+        sse = RouterSseTransport("/messages/", api_key=self.api_key)
 
         async def handle_sse(request: Request) -> None:
             async with sse.connect_sse(

src/mcpm/router/transport.py

@@ -66,8 +66,9 @@ def get_key_from_scope(scope: Scope, key_name: str) -> str | None:
 class RouterSseTransport(SseServerTransport):
     """A SSE server transport that is used by the router to handle client connections."""
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, api_key: str | None = None, **kwargs):
         self._session_id_to_identifier: dict[UUID, ClientIdentifier] = {}
+        self.api_key = api_key
         super().__init__(*args, **kwargs)
 
     @asynccontextmanager
@@ -238,6 +239,11 @@ async def handle_post_message(self, scope: Scope, receive: Receive, send: Send):
                 self._session_id_to_identifier.pop(session_id, None)
 
     def _validate_api_key(self, scope: Scope, api_key: str | None) -> bool:
+        # If we have a directly provided API key and it matches the request's API key, return True
+        if self.api_key is not None and api_key == self.api_key:
+            return True
+            
+        # Otherwise, fall back to the original validation logic
         try:
             config_manager = ConfigManager()
             host = get_key_from_scope(scope, key_name="host") or ""