doris-mcp-server/doris_mcp_server/multiworker_app.py

#!/usr/bin/env python3
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""
Multi-worker application module for doris-mcp-server

This module provides full MCP functionality with multi-worker support.
Each worker process creates its own MCP server and session manager using the same
robust architecture as the single-worker mode.
"""

import os
import asyncio
from contextlib import asynccontextmanager
import json
import logging
from typing import Any

# Import MCP components with compatibility handling
# Use the same import strategy as main.py for consistency
MCP_VERSION = 'unknown'
Server = None
InitializationOptions = None
Prompt = None
Resource = None
TextContent = None
Tool = None

def _import_mcp_with_compatibility():
    """Import MCP components with multi-version compatibility"""
    global MCP_VERSION, Server, InitializationOptions, Prompt, Resource, TextContent, Tool
    
    try:
        # Strategy 1: Try direct server-only imports to avoid client-side issues
        from mcp.server import Server as _Server
        from mcp.server.models import InitializationOptions as _InitOptions
        from mcp.types import (
            Prompt as _Prompt,
            Resource as _Resource, 
            TextContent as _TextContent,
            Tool as _Tool,
        )
        
        # Assign to globals
        Server = _Server
        InitializationOptions = _InitOptions
        Prompt = _Prompt
        Resource = _Resource
        TextContent = _TextContent
        Tool = _Tool
        
        # Try to get version safely
        try:
            import mcp
            MCP_VERSION = getattr(mcp, '__version__', None)
            if not MCP_VERSION:
                # Fallback: try to get version from package metadata
                try:
                    import importlib.metadata
                    MCP_VERSION = importlib.metadata.version('mcp')
                except Exception:
                    # Second fallback: try pkg_resources
                    try:
                        import pkg_resources
                        MCP_VERSION = pkg_resources.get_distribution('mcp').version
                    except Exception:
                        MCP_VERSION = 'detected-but-version-unknown'
        except Exception:
            # Version detection failed, but imports worked
            try:
                import importlib.metadata
                MCP_VERSION = importlib.metadata.version('mcp')
            except Exception:
                try:
                    import pkg_resources
                    MCP_VERSION = pkg_resources.get_distribution('mcp').version
                except Exception:
                    MCP_VERSION = 'imported-successfully'
            
        logger = logging.getLogger(__name__)
        logger.info(f"MCP components imported successfully in multiworker, version: {MCP_VERSION}")
        return True
        
    except Exception as import_error:
        logger = logging.getLogger(__name__)
        
        # Strategy 2: Handle RequestContext compatibility issues in 1.9.x versions
        error_str = str(import_error).lower()
        if 'requestcontext' in error_str and 'too few arguments' in error_str:
            logger.warning(f"Detected MCP RequestContext compatibility issue: {import_error}")
            logger.info("Attempting comprehensive workaround for MCP 1.9.x RequestContext issue...")
            
            try:
                # Comprehensive monkey patch approach
                import sys
                import types
                
                # Create and install mock modules before any MCP imports
                if 'mcp.shared.context' not in sys.modules:
                    mock_context_module = types.ModuleType('mcp.shared.context')
                    
                    class FlexibleRequestContext:
                        """Flexible RequestContext that accepts variable arguments"""
                        def __init__(self, *args, **kwargs):
                            self.args = args
                            self.kwargs = kwargs
                        
                        def __class_getitem__(cls, params):
                            # Accept any number of parameters and return cls
                            return cls
                        
                        # Add other methods that might be called
                        def __getattr__(self, name):
                            return lambda *args, **kwargs: None
                    
                    mock_context_module.RequestContext = FlexibleRequestContext
                    sys.modules['mcp.shared.context'] = mock_context_module
                
                # Also patch the typing system to be more permissive  
                original_check_generic = None
                try:
                    import typing
                    if hasattr(typing, '_check_generic'):
                        original_check_generic = typing._check_generic
                        def permissive_check_generic(cls, params, elen):
                            # Don't enforce strict parameter count checking
                            return
                        typing._check_generic = permissive_check_generic
                except Exception:
                    pass
                
                # Clear any cached imports that might have failed
                modules_to_clear = [k for k in sys.modules.keys() if k.startswith('mcp.')]
                for module in modules_to_clear:
                    if module in sys.modules:
                        del sys.modules[module]
                
                # Now try importing again with the patches in place
                from mcp.server import Server as _Server
                from mcp.server.models import InitializationOptions as _InitOptions
                from mcp.types import (
                    Prompt as _Prompt,
                    Resource as _Resource, 
                    TextContent as _TextContent,
                    Tool as _Tool,
                )
                
                # Assign to globals
                Server = _Server
                InitializationOptions = _InitOptions
                Prompt = _Prompt
                Resource = _Resource
                TextContent = _TextContent
                Tool = _Tool
                
                # Try to detect actual version even in compatibility mode
                try:
                    import importlib.metadata
                    actual_version = importlib.metadata.version('mcp')
                    MCP_VERSION = f'compatibility-mode-{actual_version}'
                except Exception:
                    try:
                        import pkg_resources
                        actual_version = pkg_resources.get_distribution('mcp').version
                        MCP_VERSION = f'compatibility-mode-{actual_version}'
                    except Exception:
                        MCP_VERSION = 'compatibility-mode-1.9.x'
                
                logger.info("MCP 1.9.x compatibility workaround successful in multiworker!")
                
                # Restore original typing function if we patched it
                if original_check_generic:
                    typing._check_generic = original_check_generic
                
                return True
                
            except Exception as workaround_error:
                logger.error(f"MCP compatibility workaround failed in multiworker: {workaround_error}")
                
                # Restore original typing function if we patched it
                if original_check_generic:
                    try:
                        import typing
                        typing._check_generic = original_check_generic
                    except Exception:
                        pass
        
        logger.error(f"Failed to import MCP components in multiworker: {import_error}")
        return False

# Perform MCP import with compatibility handling
if not _import_mcp_with_compatibility():
    raise ImportError(
        "Failed to import MCP components in multiworker. Please ensure MCP is properly installed. "
        "Supported versions: 1.8.x, 1.9.x"
    )

from starlette.applications import Starlette
from starlette.routing import Route
from starlette.responses import JSONResponse, Response

# Import Doris MCP components
from .tools.tools_manager import DorisToolsManager
from .tools.prompts_manager import DorisPromptsManager
from .tools.resources_manager import DorisResourcesManager
from .utils.config import DorisConfig
from .utils.db import DorisConnectionManager
from .utils.security import DorisSecurityManager

# Global variables for worker-specific instances
_worker_server = None
_worker_session_manager = None
_worker_connection_manager = None
_worker_security_manager = None
_worker_session_manager_context = None
_worker_initialized = False

def get_mcp_capabilities():
    """Get MCP capabilities for worker - use the same logic as main.py"""
    try:
        # For MCP 1.9.x and newer
        from mcp.server.lowlevel.server import NotificationOptions
        
        capabilities = {
            "resources": {},
            "tools": {},
            "prompts": {},
            "notification_options": {
                "prompts_changed": True,
                "resources_changed": True,
                "tools_changed": True
            }
        }
        return capabilities
    except Exception as e:
        # Import logger properly
        from .utils.logger import get_logger
        logger = get_logger(__name__)
        logger.warning(f"Failed to get full capabilities in multiworker: {e}")
        return {
            "resources": {},
            "tools": {},
            "prompts": {}
        }

async def initialize_worker():
    """Initialize MCP server and managers for this worker process"""
    global _worker_server, _worker_session_manager, _worker_connection_manager, _worker_security_manager, _worker_session_manager_context, _worker_initialized, _oauth_handlers, _token_handlers
    
    if _worker_initialized:
        return
    
    try:
        # Import logger properly
        from .utils.logger import get_logger
        logger = get_logger(__name__)
        
        logger.info(f"Initializing MCP worker process {os.getpid()}")
        
        # Create configuration
        config = DorisConfig.from_env()
        
        # Initialize enhanced logging system
        from .utils.config import ConfigManager
        config_manager = ConfigManager(config)
        config_manager.setup_logging()
        
        # Create security manager
        _worker_security_manager = DorisSecurityManager(config)
        
        # Initialize security manager first (includes JWT setup if enabled)
        await _worker_security_manager.initialize()
        logger.info(f"Worker {os.getpid()} security manager initialization completed")
        
        # Create connection manager with token manager for token-bound DB config
        token_manager = _worker_security_manager.auth_provider.token_manager if hasattr(_worker_security_manager, 'auth_provider') and hasattr(_worker_security_manager.auth_provider, 'token_manager') else None
        _worker_connection_manager = DorisConnectionManager(config, _worker_security_manager, token_manager)
        
        # Set connection manager reference in security manager for database validation
        _worker_security_manager.connection_manager = _worker_connection_manager
        
        await _worker_connection_manager.initialize()
        
        # Create MCP server
        _worker_server = Server("doris-mcp-server")
        
        # Create managers
        resources_manager = DorisResourcesManager(_worker_connection_manager)
        tools_manager = DorisToolsManager(_worker_connection_manager)
        prompts_manager = DorisPromptsManager(_worker_connection_manager)
        
        # Setup MCP handlers
        @_worker_server.list_resources()
        async def handle_list_resources() -> list[Resource]:
            """Handle resource list request"""
            try:
                logger.info("Handling resource list request in worker")
                resources = await resources_manager.list_resources()
                logger.info(f"Returning {len(resources)} resources from worker")
                return resources
            except Exception as e:
                logger.error(f"Failed to handle resource list request in worker: {e}")
                return []

        @_worker_server.read_resource()
        async def handle_read_resource(uri: str) -> str:
            """Handle resource read request"""
            try:
                logger.info(f"Handling resource read request in worker: {uri}")
                content = await resources_manager.read_resource(uri)
                return content
            except Exception as e:
                logger.error(f"Failed to handle resource read request in worker: {e}")
                return json.dumps(
                    {"error": f"Failed to read resource: {str(e)}", "uri": uri},
                    ensure_ascii=False,
                    indent=2,
                )

        @_worker_server.list_tools()
        async def handle_list_tools() -> list[Tool]:
            """Handle tool list request"""
            try:
                logger.info("Handling tool list request in worker")
                tools = await tools_manager.list_tools()
                logger.info(f"Returning {len(tools)} tools from worker")
                return tools
            except Exception as e:
                logger.error(f"Failed to handle tool list request in worker: {e}")
                return []

        @_worker_server.call_tool()
        async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
            """Handle tool call request"""
            try:
                logger.info(f"Handling tool call request in worker: {name}")
                result = await tools_manager.call_tool(name, arguments)
                return [TextContent(type="text", text=result)]
            except Exception as e:
                logger.error(f"Failed to handle tool call request in worker: {e}")
                error_result = json.dumps(
                    {
                        "error": f"Tool call failed: {str(e)}",
                        "tool_name": name,
                        "arguments": arguments,
                    },
                    ensure_ascii=False,
                    indent=2,
                )
                return [TextContent(type="text", text=error_result)]

        @_worker_server.list_prompts()
        async def handle_list_prompts() -> list[Prompt]:
            """Handle prompt list request"""
            try:
                logger.info("Handling prompt list request in worker")
                prompts = await prompts_manager.list_prompts()
                logger.info(f"Returning {len(prompts)} prompts from worker")
                return prompts
            except Exception as e:
                logger.error(f"Failed to handle prompt list request in worker: {e}")
                return []

        @_worker_server.get_prompt()
        async def handle_get_prompt(name: str, arguments: dict[str, Any]) -> str:
            """Handle prompt get request"""
            try:
                logger.info(f"Handling prompt get request in worker: {name}")
                result = await prompts_manager.get_prompt(name, arguments)
                return result
            except Exception as e:
                logger.error(f"Failed to handle prompt get request in worker: {e}")
                error_result = json.dumps(
                    {
                        "error": f"Failed to get prompt: {str(e)}",
                        "prompt_name": name,
                        "arguments": arguments,
                    },
                    ensure_ascii=False,
                    indent=2,
                )
                return error_result
        
        # Create session manager for this worker
        from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
        
        _worker_session_manager = StreamableHTTPSessionManager(
            app=_worker_server,
            json_response=True,
            stateless=True  # Use stateless mode for multi-worker compatibility
        )
        
        # Start the session manager context
        _worker_session_manager_context = _worker_session_manager.run()
        await _worker_session_manager_context.__aenter__()
        
        # Initialize OAuth and Token handlers
        from .auth.oauth_handlers import OAuthHandlers
        from .auth.token_handlers import TokenHandlers
        _oauth_handlers = OAuthHandlers(_worker_security_manager)
        _token_handlers = TokenHandlers(_worker_security_manager, config)
        
        _worker_initialized = True
        logger.info(f"Worker {os.getpid()} MCP initialization completed successfully")
        
    except Exception as e:
        from .utils.logger import get_logger
        logger = get_logger(__name__)
        logger.error(f"Failed to initialize worker {os.getpid()}: {e}")
        import traceback
        logger.error("Complete error stack:")
        logger.error(traceback.format_exc())
        raise

async def health_check(request):
    """Health check endpoint that shows worker PID"""
    return JSONResponse({
        "status": "healthy",
        "service": "doris-mcp-server",
        "worker_pid": os.getpid(),
        "worker_mode": "multi-process-full-mcp",
        "mcp_initialized": _worker_initialized,
        "mcp_version": MCP_VERSION
    })

# OAuth and Token handlers (initialize after worker setup)
_oauth_handlers = None
_token_handlers = None

async def oauth_login(request):
    """OAuth login endpoint"""
    if not _oauth_handlers:
        return JSONResponse({"error": "OAuth not initialized"}, status_code=503)
    return await _oauth_handlers.handle_login(request)

async def oauth_callback(request):
    """OAuth callback endpoint"""
    if not _oauth_handlers:
        return JSONResponse({"error": "OAuth not initialized"}, status_code=503)
    return await _oauth_handlers.handle_callback(request)

async def oauth_provider_info(request):
    """OAuth provider info endpoint"""
    if not _oauth_handlers:
        return JSONResponse({"error": "OAuth not initialized"}, status_code=503)
    return await _oauth_handlers.handle_provider_info(request)

async def oauth_demo(request):
    """OAuth demo page endpoint"""
    if not _oauth_handlers:
        from starlette.responses import HTMLResponse
        return HTMLResponse("<h1>OAuth not initialized</h1>")
    return await _oauth_handlers.handle_demo_page(request)

# Token management endpoints
async def token_create(request):
    """Token creation endpoint"""
    if not _token_handlers:
        return JSONResponse({"error": "Token handlers not initialized"}, status_code=503)
    return await _token_handlers.handle_create_token(request)

async def token_revoke(request):
    """Token revocation endpoint"""
    if not _token_handlers:
        return JSONResponse({"error": "Token handlers not initialized"}, status_code=503)
    return await _token_handlers.handle_revoke_token(request)

async def token_list(request):
    """Token listing endpoint"""
    if not _token_handlers:
        return JSONResponse({"error": "Token handlers not initialized"}, status_code=503)
    return await _token_handlers.handle_list_tokens(request)

async def token_stats(request):
    """Token statistics endpoint"""
    if not _token_handlers:
        return JSONResponse({"error": "Token handlers not initialized"}, status_code=503)
    return await _token_handlers.handle_token_stats(request)

async def token_cleanup(request):
    """Token cleanup endpoint"""
    if not _token_handlers:
        return JSONResponse({"error": "Token handlers not initialized"}, status_code=503)
    return await _token_handlers.handle_cleanup_tokens(request)

async def token_management(request):
    """Token management page endpoint"""
    if not _token_handlers:
        from starlette.responses import HTMLResponse
        return HTMLResponse("<h1>Token handlers not initialized</h1>")
    return await _token_handlers.handle_management_page(request)

async def root_info(request):
    """Root endpoint"""
    return JSONResponse({
        "service": "doris-mcp-server",
        "mode": "multi-worker-full-mcp",
        "worker_pid": os.getpid(),
        "mcp_initialized": _worker_initialized,
        "mcp_version": MCP_VERSION,
        "endpoints": {
            "health": "/health",
            "mcp": "/mcp"
        }
    })

@asynccontextmanager
async def lifespan(app):
    """Application lifespan manager"""
    # Startup
    try:
        await initialize_worker()
        # Import logger properly
        from .utils.logger import get_logger
        logger = get_logger(__name__)
        logger.info(f"Worker {os.getpid()} startup completed")
        
        yield
        
    finally:
        # Shutdown
        from .utils.logger import get_logger
        logger = get_logger(__name__)
        
        # Close session manager context
        if _worker_session_manager_context:
            try:
                await _worker_session_manager_context.__aexit__(None, None, None)
                logger.info(f"Worker {os.getpid()} session manager context closed")
            except Exception as e:
                logger.error(f"Error closing worker session manager context: {e}")
        
        if _worker_connection_manager:
            try:
                await _worker_connection_manager.close()
                logger.info(f"Worker {os.getpid()} connection manager closed")
            except Exception as e:
                logger.error(f"Error closing worker connection manager: {e}")
        
        if _worker_security_manager:
            try:
                await _worker_security_manager.shutdown()
                logger.info(f"Worker {os.getpid()} security manager shutdown completed")
            except Exception as e:
                logger.error(f"Error shutting down worker security manager: {e}")
        
        # Shutdown logging system
        try:
            from .utils.logger import shutdown_logging
            shutdown_logging()
        except Exception as e:
            logger.error(f"Error shutting down logging system: {e}")

async def mcp_asgi_app(scope, receive, send):
    """ASGI app that handles MCP requests"""
    if not _worker_initialized:
        # Send error response if worker not initialized
        await send({
            'type': 'http.response.start',
            'status': 503,
            'headers': [(b'content-type', b'application/json')]
        })
        await send({
            'type': 'http.response.body',
            'body': b'{"error": "Worker not initialized"}'
        })
        return
    
    # Import logger properly
    from .utils.logger import get_logger
    logger = get_logger(__name__)
    
    # Get request path for logging
    path = scope.get('path', '')
    method = scope.get('method', 'UNKNOWN')
    logger.debug(f"Worker {os.getpid()} handling MCP request: {method} {path}")
    
    # Handle the request directly without nested run context
    await _worker_session_manager.handle_request(scope, receive, send)

# Create Starlette app with basic routes
basic_app = Starlette(
    debug=True,
    routes=[
        Route("/", root_info, methods=["GET"]),
        Route("/health", health_check, methods=["GET"]),
        # OAuth endpoints
        Route("/auth/login", oauth_login, methods=["GET"]),
        Route("/auth/callback", oauth_callback, methods=["GET"]),
        Route("/auth/provider", oauth_provider_info, methods=["GET"]),
        Route("/auth/demo", oauth_demo, methods=["GET"]),
        # Token management endpoints
        Route("/token/create", token_create, methods=["GET", "POST"]),
        Route("/token/revoke", token_revoke, methods=["GET", "DELETE"]),
        Route("/token/list", token_list, methods=["GET"]),
        Route("/token/stats", token_stats, methods=["GET"]),
        Route("/token/cleanup", token_cleanup, methods=["GET", "POST"]),
        Route("/token/management", token_management, methods=["GET"]),
    ],
    lifespan=lifespan
)

# Create main ASGI app that routes between basic app and MCP
async def app(scope, receive, send):
    """Main ASGI app that routes requests"""
    path = scope.get('path', '/')
    
    if path == "/mcp" or path.startswith('/mcp/'):
        # Handle MCP requests with session manager
        await mcp_asgi_app(scope, receive, send)
    else:
        # Handle other requests with basic Starlette app (includes auth endpoints)
        await basic_app(scope, receive, send)