MCP Server Hero

"""Logging middleware for request/response logging.""" import time from typing import Any, Callable, Dict, List import mcp.types as types from .base import BaseMiddleware from ..utils.logging import get_logger logger = get_logger(__name__) class LoggingMiddleware(BaseMiddleware): """Middleware for comprehensive request/response logging.""" def __init__(self, log_level: str = "INFO", log_arguments: bool = True, log_results: bool = False) -> None: """Initialize logging middleware. Args: log_level: Logging level (DEBUG, INFO, WARNING, ERROR) log_arguments: Whether to log arguments log_results: Whether to log results (be careful with sensitive data) """ super().__init__("logging") self.log_level = log_level self.log_arguments = log_arguments self.log_results = log_results self._logger = get_logger(f"{__name__}.LoggingMiddleware") async def process_tool_call( self, name: str, arguments: Dict[str, Any], next_handler: Callable[[str, Dict[str, Any]], Any] ) -> List[types.ContentBlock]: """Process tool call with logging.""" start_time = time.time() # Log request log_msg = f"Tool call: {name}" if self.log_arguments: log_msg += f" with arguments: {arguments}" self._logger.info(log_msg) try: result = await next_handler(name, arguments) # Log success duration = time.time() - start_time success_msg = f"Tool {name} completed in {duration:.3f}s" if self.log_results: success_msg += f" with result: {result}" self._logger.info(success_msg) return result except Exception as e: # Log error duration = time.time() - start_time self._logger.error(f"Tool {name} failed after {duration:.3f}s: {str(e)}") raise async def process_resource_read(self, uri: Any, next_handler: Callable[[Any], Any]) -> str: """Process resource read with logging.""" start_time = time.time() # Log request self._logger.info(f"Resource read: {uri}") try: result = await next_handler(uri) # Log success duration = time.time() - start_time success_msg = f"Resource {uri} read in {duration:.3f}s" if self.log_results: content_preview = result[:100] + "..." if len(result) > 100 else result success_msg += f" (preview: {content_preview})" self._logger.info(success_msg) return result except Exception as e: # Log error duration = time.time() - start_time self._logger.error(f"Resource {uri} read failed after {duration:.3f}s: {str(e)}") raise async def process_prompt_get( self, name: str, arguments: Dict[str, str] | None, next_handler: Callable[[str, Dict[str, str] | None], Any] ) -> types.GetPromptResult: """Process prompt get with logging.""" start_time = time.time() # Log request log_msg = f"Prompt get: {name}" if self.log_arguments and arguments: log_msg += f" with arguments: {arguments}" self._logger.info(log_msg) try: result = await next_handler(name, arguments) # Log success duration = time.time() - start_time success_msg = f"Prompt {name} generated in {duration:.3f}s" if self.log_results: msg_count = len(result.messages) if result.messages else 0 success_msg += f" ({msg_count} messages)" self._logger.info(success_msg) return result except Exception as e: # Log error duration = time.time() - start_time self._logger.error(f"Prompt {name} generation failed after {duration:.3f}s: {str(e)}") raise