gnosys-strata 1.1.4__py3-none-any.whl

strata/mcp_client_manager.py

@@ -0,0 +1,282 @@
+"""MCP Client Manager for managing multiple MCP server connections."""
+
+import asyncio
+import logging
+from pathlib import Path
+from typing import Dict, List, Optional
+
+from strata.config import MCPServerConfig, MCPServerList
+from strata.mcp_proxy.client import MCPClient
+from strata.mcp_proxy.transport.http import HTTPTransport
+from strata.mcp_proxy.transport.stdio import StdioTransport
+from strata.utils.catalog import ToolCatalog
+
+logger = logging.getLogger(__name__)
+
+
+class MCPClientManager:
+    """Manages multiple MCP client connections based on configuration."""
+
+    def __init__(self, config_path: Optional[Path] = None, server_names: Optional[List[str]] = None):
+        """Initialize the MCP client manager.
+
+        Args:
+            config_path: Optional path to configuration file.
+                        If None, uses default from MCPServerList.
+            server_names: Optional list of specific server names to initialize.
+                         If None, all enabled servers will be initialized.
+        """
+        self.server_list = MCPServerList(config_path)
+        self.server_names = server_names  # Specific servers to manage
+        self.active_clients: Dict[str, MCPClient] = {}
+        self.active_transports: Dict[str, HTTPTransport | StdioTransport] = {}
+        # Cache of current server configs for comparison during sync
+        self.cached_configs: List[MCPServerConfig] = []
+        self.catalog = ToolCatalog()
+        # Mutex to prevent concurrent sync operations
+        self._sync_lock = asyncio.Lock()
+
+    async def initialize_from_config(self) -> Dict[str, bool]:
+        """Initialize MCP clients from configuration.
+
+        Only initializes servers that are enabled in the configuration.
+        If server_names was specified in __init__, only those servers will be initialized.
+
+        Returns:
+            Dict mapping server names to success status (True if connected)
+        """
+        results = {}
+        enabled_servers = self.server_list.list_servers(enabled_only=True)
+        
+        # Filter servers if specific names were provided
+        if self.server_names:
+            enabled_servers = [s for s in enabled_servers if s.name in self.server_names]
+
+        for server in enabled_servers:
+            try:
+                await self._connect_server(server)
+                results[server.name] = True
+                logger.info(f"Successfully connected to MCP server: {server.name}")
+            except Exception as e:
+                results[server.name] = False
+                logger.error(f"Failed to connect to MCP server {server.name}: {e}")
+
+        # Cache all server configs (both enabled and disabled) for future comparisons
+        self.cached_configs = self.server_list.list_servers()
+
+        return results
+    
+    async def authenticate_server(self, server_name: str) -> None:
+        """Authenticate a single MCP server.
+
+        Args:
+            server_name: Name of the server to authenticate
+        """
+        if server_name in self.active_clients:
+            client = self.active_clients[server_name]
+            try:
+                await client.initialize()
+                await client.disconnect()
+                logger.info(f"Initialized MCP server: {server_name}")
+            except Exception as e:
+                logger.error(f"Error initializing {server_name}: {e}")
+
+    async def _connect_server(self, server: MCPServerConfig) -> None:
+        """Connect to a single MCP server.
+
+        Args:
+            server: Server configuration
+
+        Raises:
+            Exception: If connection fails
+        """
+        # Create transport based on type
+        if server.type in ["sse", "http"]:
+            if not server.url:
+                raise ValueError(f"Server {server.name} has no URL configured")
+
+            transport = HTTPTransport(
+                server_name=server.name,
+                url=server.url,
+                mode=server.type,  # "http" or "sse" # type: ignore
+                headers=server.headers,
+                auth=server.auth
+            )
+        else:  # stdio/command
+            if not server.command:
+                raise ValueError(f"Server {server.name} has no command configured")
+            transport = StdioTransport(
+                command=server.command, args=server.args, env=server.env
+            )
+
+        # Create client
+        client = MCPClient(transport)
+
+        # Connect
+        await client.connect()
+
+        # Store active client and transport
+        self.active_clients[server.name] = client
+        self.active_transports[server.name] = transport
+
+        # Auto-populate catalog on first connect (background task, non-blocking)
+        if not self.catalog.get_tools(server.name):
+            async def populate_catalog():
+                try:
+                    tools = await client.list_tools()
+                    self.catalog.update_server(server.name, tools)
+                    logger.info(f"Added {server.name} to catalog (first connection)")
+                except Exception as e:
+                    logger.warning(f"Could not update catalog for {server.name}: {e}")
+            asyncio.create_task(populate_catalog())
+
+    async def _disconnect_server(self, server_name: str) -> None:
+        """Disconnect from a single MCP server.
+
+        Args:
+            server_name: Name of the server to disconnect
+        """
+        if server_name in self.active_clients:
+            client = self.active_clients[server_name]
+            try:
+                await client.disconnect()
+                logger.info(f"Disconnected from MCP server: {server_name}")
+            except Exception as e:
+                logger.error(f"Error disconnecting from {server_name}: {e}")
+            finally:
+                # Remove from active clients and transports
+                if server_name in self.active_clients:
+                    del self.active_clients[server_name]
+                if server_name in self.active_transports:
+                    del self.active_transports[server_name]
+
+    async def sync_with_config(self, new_servers: Dict[str, MCPServerConfig]) -> None:
+        """Sync the manager state with new configuration.
+
+        This method handles all changes: add, remove, enable, disable, and config updates.
+        Uses a mutex lock to prevent concurrent sync operations.
+
+        Args:
+            new_servers: New server configurations from config file
+        """
+        async with self._sync_lock:
+            # Create lookup for current cached configs by name
+            cached_by_name = {config.name: config for config in self.cached_configs}
+
+            # Find servers to remove (in active clients but not in new config)
+            servers_to_remove = set(self.active_clients.keys()) - set(
+                new_servers.keys()
+            )
+            for server_name in servers_to_remove:
+                await self._disconnect_server(server_name)
+                logger.info(f"Removed MCP server: {server_name}")
+
+            # Process each server in new config
+            for server_name, new_config in new_servers.items():
+                try:
+                    is_active = server_name in self.active_clients
+                    cached_config = cached_by_name.get(server_name)
+                    config_changed = cached_config != new_config
+
+                    if new_config.enabled:
+                        if not is_active:
+                            # Server is enabled but not connected, connect it
+                            await self._connect_server(new_config)
+                            logger.info(f"Connected to MCP server: {server_name}")
+                        elif config_changed:
+                            # Server is active but config changed, reconnect
+                            await self._disconnect_server(server_name)
+                            await self._connect_server(new_config)
+                            logger.info(
+                                f"Reconnected MCP server with new config: {server_name}"
+                            )
+                        # If server is active and config unchanged, do nothing
+                    else:
+                        if is_active:
+                            # Server is disabled but still connected, disconnect it
+                            await self._disconnect_server(server_name)
+                            logger.info(f"Disabled MCP server: {server_name}")
+                except Exception as e:
+                    logger.error(f"Failed to sync server {server_name}: {e}")
+
+            # Update cached configs with new config
+            self.cached_configs = list(new_servers.values())
+
+    def get_client(self, server_name: str) -> MCPClient:
+        """Get an active MCP client by server name.
+
+        Args:
+            server_name: Name of the server
+
+        Returns:
+            MCPClient instance if active, None otherwise
+        """
+        return self.active_clients[server_name]
+
+    def list_active_servers(self) -> list[str]:
+        """List names of all active (connected) servers.
+
+        Returns:
+            List of active server names
+        """
+        return list(self.active_clients.keys())
+
+    def is_connected(self, server_name: str) -> bool:
+        """Check if a server is currently connected.
+
+        Args:
+            server_name: Name of the server
+
+        Returns:
+            True if connected, False otherwise
+        """
+        client = self.active_clients.get(server_name)
+        return client is not None and client.is_connected()
+
+    async def disconnect_all(self) -> None:
+        """Disconnect from all active MCP servers."""
+        server_names = list(self.active_clients.keys())
+        for server_name in server_names:
+            await self._disconnect_server(server_name)
+        logger.info("Disconnected from all MCP servers")
+
+    async def reconnect_server(self, server_name: str) -> bool:
+        """Reconnect to a server (disconnect if connected, then connect).
+
+        Args:
+            server_name: Name of the server to reconnect
+
+        Returns:
+            True if successfully reconnected, False otherwise
+        """
+        server = self.server_list.get_server(server_name)
+        if server is None:
+            logger.error(f"Server not found: {server_name}")
+            return False
+
+        if not server.enabled:
+            logger.error(f"Cannot reconnect disabled server: {server_name}")
+            return False
+
+        try:
+            # Disconnect if active
+            if server_name in self.active_clients:
+                await self._disconnect_server(server_name)
+
+            # Reconnect
+            await self._connect_server(server)
+            logger.info(f"Reconnected to MCP server: {server_name}")
+            return True
+
+        except Exception as e:
+            logger.error(f"Failed to reconnect to server {server_name}: {e}")
+            return False
+
+    async def __aenter__(self):
+        """Enter async context manager."""
+        await self.initialize_from_config()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Exit async context manager."""
+        await self.disconnect_all()

strata/mcp_proxy/__init__.py

@@ -0,0 +1,7 @@
+"""MCP Proxy module for connecting to and interacting with MCP servers."""
+
+from .client import MCPClient
+from .transport import HTTPTransport, StdioTransport, Transport
+from .auth_provider import create_oauth_provider
+
+__all__ = ["MCPClient", "StdioTransport", "HTTPTransport", "Transport", "create_oauth_provider"]

strata/mcp_proxy/auth_provider.py

@@ -0,0 +1,200 @@
+import threading
+import time
+import webbrowser
+import os
+import json
+import logging
+
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from urllib.parse import parse_qs, urlparse
+
+from mcp.client.auth import OAuthClientProvider, TokenStorage
+from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthToken
+
+
+# Set up logging
+logger = logging.getLogger(__name__)
+
+TOKEN_DIRECTORY = '.tokens'
+
+class LocalTokenStorage(TokenStorage):
+    """Simple in-memory token storage implementation."""
+
+    def __init__(self, server_name: str = "default"):
+        self.server_name: str = server_name
+        self._tokens: OAuthToken | None = None
+        self._client_info: OAuthClientInformationFull | None = None
+        self.token_lock = threading.Lock()
+        self.info_lock = threading.Lock()
+        self.TOKEN_PATH = os.path.join(TOKEN_DIRECTORY, self.server_name,"tokens.json")
+
+    async def get_tokens(self) -> OAuthToken | None:
+        if os.path.exists(self.TOKEN_PATH) and self._tokens is None:
+            with self.token_lock:
+                with open(self.TOKEN_PATH, "r") as f:
+                    try:
+                        data = json.load(f)
+                        self._tokens = OAuthToken.model_validate(data)
+                    except Exception as e:
+                        logger.info("Error loading tokens:", e)
+        return self._tokens
+
+    async def set_tokens(self, tokens: OAuthToken) -> None:
+        self._tokens = tokens
+        os.makedirs(os.path.dirname(self.TOKEN_PATH), exist_ok=True)
+        with self.token_lock:
+            with open(self.TOKEN_PATH, "w") as f:
+                dump = tokens.model_dump(exclude_none=True, mode='json')
+                json.dump(dump, f)
+
+    async def get_client_info(self) -> OAuthClientInformationFull | None:
+        return self._client_info
+
+    async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
+        self._client_info = client_info
+
+
+class CallbackHandler(BaseHTTPRequestHandler):
+    """Simple HTTP handler to capture OAuth callback."""
+
+    def __init__(self, request, client_address, server, callback_data):
+        """Initialize with callback data storage."""
+        self.callback_data = callback_data
+        super().__init__(request, client_address, server)
+
+    def do_GET(self):
+        """Handle GET request from OAuth redirect."""
+        parsed = urlparse(self.path)
+        query_params = parse_qs(parsed.query)
+
+        if "code" in query_params:
+            self.callback_data["authorization_code"] = query_params["code"][0]
+            self.callback_data["state"] = query_params.get("state", [None])[0]
+            self.send_response(200)
+            self.send_header("Content-type", "text/html")
+            self.end_headers()
+            self.wfile.write(b"""
+            <html>
+            <body>
+                <h1>Authorization Successful!</h1>
+                <p>You can close this window and return to the terminal.</p>
+                <script>setTimeout(() => window.close(), 2000);</script>
+            </body>
+            </html>
+            """)
+        elif "error" in query_params:
+            self.callback_data["error"] = query_params["error"][0]
+            self.send_response(400)
+            self.send_header("Content-type", "text/html")
+            self.end_headers()
+            self.wfile.write(
+                f"""
+            <html>
+            <body>
+                <h1>Authorization Failed</h1>
+                <p>Error: {query_params["error"][0]}</p>
+                <p>You can close this window and return to the terminal.</p>
+            </body>
+            </html>
+            """.encode()
+            )
+        else:
+            self.send_response(404)
+            self.end_headers()
+
+    def log_message(self, format, *args):
+        """Suppress default logging."""
+        pass
+
+
+class CallbackServer:
+    """Simple server to handle OAuth callbacks."""
+
+    def __init__(self, port=3000):
+        self.port = port
+        self.server = None
+        self.thread = None
+        self.callback_data = {"authorization_code": None, "state": None, "error": None}
+
+    def _create_handler_with_data(self):
+        """Create a handler class with access to callback data."""
+        callback_data = self.callback_data
+
+        class DataCallbackHandler(CallbackHandler):
+            def __init__(self, request, client_address, server):
+                super().__init__(request, client_address, server, callback_data)
+
+        return DataCallbackHandler
+
+    def start(self):
+        """Start the callback server in a background thread."""
+        handler_class = self._create_handler_with_data()
+        self.server = HTTPServer(("localhost", self.port), handler_class)
+        self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
+        self.thread.start()
+        logger.info(f"🖥️  Started callback server on http://localhost:{self.port}")
+
+    def stop(self):
+        """Stop the callback server."""
+        if self.server:
+            self.server.shutdown()
+            self.server.server_close()
+        if self.thread:
+            self.thread.join(timeout=1)
+
+    def wait_for_callback(self, timeout=300):
+        """Wait for OAuth callback with timeout."""
+        start_time = time.time()
+        while time.time() - start_time < timeout:
+            if self.callback_data["authorization_code"]:
+                return self.callback_data["authorization_code"]
+            elif self.callback_data["error"]:
+                raise Exception(f"OAuth error: {self.callback_data['error']}")
+            time.sleep(0.1)
+        raise Exception("Timeout waiting for OAuth callback")
+
+    def get_state(self):
+        """Get the received state parameter."""
+        return self.callback_data["state"]
+    
+def create_oauth_provider(server_name: str, url: str) -> OAuthClientProvider:
+    """Create OAuth authentication provider."""
+    client_metadata_dict = {
+        "client_name": "KLAVIS Strata MCP Router",
+        "redirect_uris": ["http://localhost:3030/callback"],
+        "grant_types": ["authorization_code", "refresh_token"],
+        "response_types": ["code"],
+        "token_endpoint_auth_method": "client_secret_post",
+    }
+
+    async def callback_handler() -> tuple[str, str | None]:
+        """Wait for OAuth callback and return auth code and state."""
+        callback_server = CallbackServer(port=3030)
+        callback_server.start()
+        logger.info("⏳ Waiting for authorization callback...")
+        try:
+            auth_code = callback_server.wait_for_callback(timeout=300)
+            return auth_code, callback_server.get_state()
+        finally:
+            callback_server.stop()
+
+    async def _default_redirect_handler(authorization_url: str) -> None:
+        """Default redirect handler that opens the URL in a browser."""
+        logger.info(f"Opening browser for authorization: {authorization_url}")
+        webbrowser.open(authorization_url)
+
+    auth_url = ""
+    if url.endswith("/mcp"):
+        auth_url = url[:-4]
+    elif url.endswith("/sse"):
+        auth_url = url[:-4]
+    else:
+        auth_url = url
+    
+    return OAuthClientProvider(
+        server_url=auth_url,
+        client_metadata=OAuthClientMetadata.model_validate(client_metadata_dict),
+        storage=LocalTokenStorage(server_name),
+        redirect_handler=_default_redirect_handler,
+        callback_handler=callback_handler,
+    )

strata/mcp_proxy/client.py

@@ -0,0 +1,162 @@
+"""MCP Client for connecting to and interacting with MCP servers."""
+
+import logging
+from typing import Any, Dict, List, Optional
+
+from mcp import types
+
+from .transport import Transport
+
+logger = logging.getLogger(__name__)
+
+
+class MCPClient:
+    """Client for connecting to MCP servers using various transports.
+
+    Usage:
+        # With stdio transport
+        transport = StdioTransport("docker", ["run", "-i", "my-server"])
+        client = MCPClient(transport)
+        await client.connect()
+
+        # With HTTP/SSE transport
+        transport = HTTPTransport("http://localhost:8080", mode="sse")
+        client = MCPClient(transport)
+        await client.connect()
+    """
+
+    def __init__(self, transport: Transport):
+        """Initialize the MCP client with a transport.
+
+        Args:
+            transport: Transport instance (StdioTransport or HTTPTransport)
+        """
+        self.transport = transport
+        self._tools_cache: Optional[List[Dict[str, Any]]] = None
+
+    async def initialize(self) -> None:
+        """Initialize the MCP client by connecting the transport."""
+        await self.transport.initialize()
+
+    async def connect(self) -> None:
+        """Connect to the MCP server."""
+        await self.transport.connect()
+        logger.info(
+            f"Connected to MCP server using {self.transport.__class__.__name__}"
+        )
+
+    async def disconnect(self) -> None:
+        """Disconnect from the MCP server."""
+        await self.transport.disconnect()
+        self._tools_cache = None
+        logger.info("Disconnected from MCP server")
+
+    def is_connected(self) -> bool:
+        """Check if connected to an MCP server."""
+        return self.transport.is_connected()
+
+    async def list_tools(self, use_cache: bool = True) -> List[Dict[str, Any]]:
+        """List available tools from the MCP server.
+
+        Args:
+            use_cache: Whether to use cached tools if available
+
+        Returns:
+            List of tool definitions with name, description, and inputSchema
+        """
+        if not self.transport.is_connected():
+            raise RuntimeError("Not connected to any MCP server")
+
+        if use_cache and self._tools_cache is not None:
+            return self._tools_cache
+
+        # Get tools from server
+        session = self.transport.get_session()
+        response = await session.list_tools()
+
+        # Convert to dict format
+        tools = []
+        for tool in response.tools:
+            tool_dict = {
+                "name": tool.name,
+                "description": tool.description,
+                "inputSchema": tool.inputSchema,
+            }
+            # Add optional fields only if they have values
+            if hasattr(tool, "title") and tool.title:
+                tool_dict["title"] = tool.title
+            if hasattr(tool, "outputSchema") and tool.outputSchema:
+                tool_dict["outputSchema"] = tool.outputSchema
+            tools.append(tool_dict)
+
+        self._tools_cache = tools
+        logger.info(f"Retrieved {len(tools)} tools from MCP server")
+
+        return tools
+
+    async def call_tool(
+        self, tool_name: str, arguments: Dict[str, Any]
+    ) -> List[types.ContentBlock]:
+        """Call a tool on the MCP server.
+
+        Args:
+            tool_name: Name of the tool to call
+            arguments: Arguments to pass to the tool
+
+        Returns:
+            Tool execution result from MCP server
+        """
+        if not self.transport.is_connected():
+            raise RuntimeError("Not connected to any MCP server")
+
+        logger.info(f"Calling tool '{tool_name}' with arguments: {arguments}")
+
+        # Call the tool and return result directly
+        session = self.transport.get_session()
+        result = await session.call_tool(tool_name, arguments)
+        if result.isError:
+            # Extract error text from content blocks (stderr is captured here)
+            logger.debug(f"Error result content: {result.content}")
+            logger.debug(f"Error result content type: {type(result.content)}")
+            logger.debug(f"Error result content length: {len(result.content) if result.content else 0}")
+
+            error_messages = []
+            for i, content_block in enumerate(result.content or []):
+                logger.debug(f"Content block {i}: {content_block}, type: {type(content_block)}")
+                if hasattr(content_block, 'text'):
+                    logger.debug(f"Content block {i} has text: {content_block.text}")
+                    error_messages.append(content_block.text)
+                else:
+                    logger.debug(f"Content block {i} attributes: {dir(content_block)}")
+
+            error_text = "\n".join(error_messages) if error_messages else f"Unknown error (no error text in {len(result.content or [])} content blocks)"
+
+            logger.error(f"Tool '{tool_name}' returned error: {error_text}")
+            raise RuntimeError(f"Tool '{tool_name}' error: {error_text}")
+        return result.content
+
+    async def get_tool_schema(self, tool_name: str) -> Optional[Dict[str, Any]]:
+        """Get the schema for a specific tool.
+
+        Args:
+            tool_name: Name of the tool
+
+        Returns:
+            Tool schema or None if not found
+        """
+        tools = await self.list_tools()
+
+        for tool in tools:
+            if tool["name"] == tool_name:
+                return tool
+
+        return None
+
+    async def __aenter__(self):
+        """Enter async context manager."""
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Exit async context manager."""
+        await self.disconnect()

strata/mcp_proxy/transport/__init__.py

@@ -0,0 +1,7 @@
+"""Transport implementations for MCP client."""
+
+from .base import Transport
+from .http import HTTPTransport
+from .stdio import StdioTransport
+
+__all__ = ["Transport", "HTTPTransport", "StdioTransport"]