#!/usr/bin/env python3 """ Enhanced MCP server with dynamic tool update support. This server extends the ManualMCPServer to support real-time tool updates through WebSocket, SSE, and polling mechanisms. """ import asyncio import json import logging import os import sys import threading from datetime import datetime from typing import Any, Dict, List, Optional, Set from .base_server import BaseMCPServer from .update_listener import UpdateListener, UpdateListenerConfig from .startup_fix import ensure_gateway_ready logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class DynamicMCPServer(BaseMCPServer): """ Enhanced MCP server that supports dynamic tool updates. Features: - Real-time tool updates via WebSocket/SSE/polling - Graceful connection handling and reconnection - Tool change notifications to clients - Backward compatibility with standard MCP protocol """ def __init__(self): """Initialize the dynamic MCP server.""" # Ensure gateway is ready before initialization try: ensure_gateway_ready() except RuntimeError as e: logger.error(f"Failed to connect to gateway: {e}") # Continue with limited functionality super().__init__() # Dynamic capabilities self.supports_dynamic_updates = True self.update_listener: Optional[UpdateListener] = None self.update_thread: Optional[threading.Thread] = None self.update_event = threading.Event() # Track tool versions for change detection self.tool_versions: Dict[str, int] = {} self._tools_lock = threading.Lock() # Pending notifications self.pending_notifications: List[Dict[str, Any]] = [] self._notifications_lock = threading.Lock() # Initialize update listener self._initialize_update_listener() def _initialize_update_listener(self): """Initialize the update listener based on configuration.""" try: # Get configuration from environment gateway_url = os.getenv("GATEWAY_URL", "http://localhost:8000") update_method = os.getenv("MCP_UPDATE_METHOD", "websocket") # websocket, sse, or polling polling_interval = int(os.getenv("MCP_POLLING_INTERVAL", "30")) # Create listener configuration config = UpdateListenerConfig( gateway_url=gateway_url, bundle_name=self.bundle_name, update_method=update_method, polling_interval=polling_interval, on_tools_changed=self._handle_tools_changed, on_connection_error=self._handle_connection_error ) # Create and start listener self.update_listener = UpdateListener(config) self.update_thread = threading.Thread( target=self._run_update_listener, daemon=True, name="MCP-UpdateListener" ) self.update_thread.start() logger.info(f"Initialized update listener with method: {update_method}") except Exception as e: logger.error(f"Failed to initialize update listener: {e}") # Continue without dynamic updates self.supports_dynamic_updates = False def _run_update_listener(self): """Run the update listener in a separate thread.""" try: # Create event loop for this thread loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) # Run the listener loop.run_until_complete(self.update_listener.start()) except Exception as e: logger.error(f"Update listener error: {e}") self.supports_dynamic_updates = False finally: # Clean up the event loop try: loop.close() except Exception: pass def _handle_tools_changed(self, event: Dict[str, Any]): """ Handle tool change events from the update listener. Args: event: Tool change event data """ try: event_type = event.get("event_type") tool_id = event.get("tool_id") tool_name = event.get("tool_name") logger.info(f"Received tool change event: {event_type} for {tool_name}") # Reload tools from gateway with self._tools_lock: old_tools = self.tools.copy() self.load_tools() new_tools = self.tools # Detect changes old_tool_names = {t["name"] for t in old_tools} new_tool_names = {t["name"] for t in new_tools} added = new_tool_names - old_tool_names removed = old_tool_names - new_tool_names # Check for updates (same name but different content) updated = set() for tool in new_tools: name = tool["name"] if name in old_tool_names: old_tool = next(t for t in old_tools if t["name"] == name) if tool != old_tool: updated.add(name) # Update tool versions for tool in new_tools: name = tool["name"] self.tool_versions[name] = self.tool_versions.get(name, 0) + 1 # Queue notification if there were changes if added or removed or updated: self._queue_tools_changed_notification({ "added": list(added), "removed": list(removed), "updated": list(updated), "total_tools": len(new_tools) }) except Exception as e: logger.error(f"Error handling tool change: {e}") def _handle_connection_error(self, error: Exception): """ Handle connection errors from the update listener. Args: error: The connection error """ logger.warning(f"Update listener connection error: {error}") # The listener will handle reconnection automatically def _queue_tools_changed_notification(self, changes: Dict[str, Any]): """ Queue a tools changed notification to be sent to the client. Args: changes: Dictionary describing the changes """ with self._notifications_lock: notification = { "jsonrpc": "2.0", "method": "notifications/tools/list_changed", "params": { "timestamp": datetime.utcnow().isoformat(), "changes": changes } } self.pending_notifications.append(notification) # Signal that we have notifications self.update_event.set() def _send_pending_notifications(self): """Send any pending notifications to the client.""" with self._notifications_lock: for notification in self.pending_notifications: try: sys.stdout.write(json.dumps(notification) + "\n") sys.stdout.flush() logger.info(f"Sent notification: {notification['method']}") except Exception as e: logger.error(f"Failed to send notification: {e}") self.pending_notifications.clear() def handle_initialize(self, request_id): """ Handle initialize request with dynamic update capabilities. Override to indicate support for dynamic tool updates. """ response = super().handle_initialize(request_id) # Update capabilities to indicate dynamic tool support if self.supports_dynamic_updates: response["result"]["capabilities"]["tools"]["listChanged"] = True # Add experimental capabilities (for future extensions) response["result"]["capabilities"]["experimental"] = { "dynamicTools": True, "updateMethods": ["websocket", "sse", "polling"], "bundleSupport": True } return response def handle_tools_list(self, request_id): """ Handle tools/list request with version tracking. Override to include tool versions for change detection. """ response = super().handle_tools_list(request_id) # Add version information if dynamic updates are supported if self.supports_dynamic_updates: with self._tools_lock: for tool in response["result"]["tools"]: tool_name = tool["name"] tool["_version"] = self.tool_versions.get(tool_name, 1) return response def run(self): """ Run the dynamic MCP server. Override to handle notifications in addition to requests. """ logger.info(f"Starting Dynamic MCP Server (bundle: {self.bundle_name or 'all'})") logger.info(f"Tools loaded: {len(self.tools)}") logger.info(f"Dynamic updates supported: {self.supports_dynamic_updates}") logger.info(f"Running in stdin mode, waiting for messages...") # Flush stderr to ensure logs are visible sys.stderr.flush() # Use select or similar for non-blocking I/O import select import time # Track if we've received any messages received_messages = False startup_time = time.time() consecutive_eof_count = 0 # In Docker/containerized environments, stdin might not be immediately available # Give it some time before considering EOF as a disconnect STARTUP_GRACE_PERIOD = 60 # seconds # Check if we're in a containerized environment is_docker = os.environ.get('MCP_MODE') == 'docker' while True: try: # Check for pending notifications if self.update_event.is_set(): self._send_pending_notifications() self.update_event.clear() # Check if stdin is available for reading try: # In some environments, select might fail on stdin readable, _, _ = select.select([sys.stdin], [], [], 0.1) except (ValueError, OSError) as e: # stdin might not support select in some environments logger.debug(f"select() failed on stdin: {e}, falling back to blocking read") readable = [sys.stdin] if readable: try: line = sys.stdin.readline() if not line: # EOF detected consecutive_eof_count += 1 elapsed = time.time() - startup_time # In Docker mode without any client messages, just wait if is_docker and not received_messages: if consecutive_eof_count == 1: logger.info("MCP server ready, waiting for client connection...") elif consecutive_eof_count % 30 == 0: # Log every ~30 seconds logger.debug(f"Still waiting for client (elapsed: {elapsed:.0f}s)") time.sleep(1) # Wait longer between checks continue # Otherwise check startup grace period if elapsed < STARTUP_GRACE_PERIOD and not received_messages: # During startup grace period, ignore EOF logger.debug(f"Ignoring EOF during startup (elapsed: {elapsed:.1f}s)") time.sleep(0.5) # Avoid busy loop continue else: # After grace period or after receiving messages logger.info("Client disconnected (EOF on stdin)") break else: # Reset EOF counter on successful read consecutive_eof_count = 0 except IOError: # stdin might be closed in some container environments logger.debug("IOError reading stdin, continuing...") time.sleep(0.5) continue # If we got a line with content, process it if line and line.strip(): request = json.loads(line.strip()) method = request.get("method") request_id = request.get("id") # Mark that we've received a message received_messages = True # Handle request as before if method == "initialize": response = self.handle_initialize(request_id) elif method == "tools/list": response = self.handle_tools_list(request_id) elif method == "tools/call": response = self.handle_tools_call(request_id, request.get("params", {})) elif method == "resources/list": response = self.handle_resources_list(request_id) elif method == "prompts/list": response = self.handle_prompts_list(request_id) elif method == "notifications/initialized": # This is a notification, no response needed continue elif method == "notifications/tools/list_changed": # Client acknowledging our notification continue else: response = { "jsonrpc": "2.0", "id": request_id, "error": {"code": -32601, "message": "Method not found"} } # Send response sys.stdout.write(json.dumps(response) + "\n") sys.stdout.flush() except json.JSONDecodeError as e: logger.error(f"Invalid JSON received: {e}") # Don't send error response for JSON decode errors as we don't have request_id except KeyboardInterrupt: logger.info("Received interrupt, shutting down...") break except Exception as e: logger.error(f"Server error: {e}", exc_info=True) if 'request' in locals() and 'id' in request: error_response = { "jsonrpc": "2.0", "id": request.get("id"), "error": {"code": -32603, "message": str(e)} } sys.stdout.write(json.dumps(error_response) + "\n") sys.stdout.flush() # Cleanup if self.update_listener: # Create a new event loop for cleanup if needed try: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) loop.run_until_complete(self.update_listener.stop()) loop.close() except Exception as e: logger.error(f"Error during cleanup: {e}") logger.info("Dynamic MCP Server stopped") if __name__ == "__main__": server = DynamicMCPServer() server.run()