MCPX MCP Gateway

import asyncio from tornado.log import gen_log as logger from json import loads from typing import Dict from tornado.websocket import WebSocketHandler DISCOVERY_EVENT = "discovery-event" METRICS_EVENT = "instance-level-metrics" CONFIGURATION_LOAD_EVENT = "configuration-load-event" EVENT_KEY = "event" DATA_KEY = "data" CONNECTION_READY_MESSAGE = "ready" class WebSocketHubHandler(WebSocketHandler): def initialize(self, cache: Dict[str, Dict[str, str]], lock: asyncio.Lock): self.cache = cache self.lock = lock async def on_message(self, message: str | bytes): logger.info(f"Received message: {message}") await self.handle_message(message) async def handle_message(self, message: str | bytes): try: json_data = loads(message) logger.info(f"Parsed message: {json_data}") event_name = json_data.get(EVENT_KEY) logger.info(f"Event: {event_name}") handler = self.get_event_handler(event_name) if handler is None: print(f"Error handling event message of type: {event_name}") else: await handler(json_data) except Exception as e: print(f"Error handling message: {e}") def get_event_handler(self, event: str): logger.info(f"Event: {event}") return { DISCOVERY_EVENT: self._discovery_event, CONFIGURATION_LOAD_EVENT: self._configuration_load_event, METRICS_EVENT: self._metrics_event, }.get(event) async def _discovery_event(self, event_data: Dict[str, str]): async with self.lock: logger.info(f"Discovery Event: {event_data}") self.cache[DISCOVERY_EVENT] = event_data async def _configuration_load_event(self, event_data: Dict[str, str]): async with self.lock: logger.info(f"Configuration Load Event: {event_data}") self.cache[CONFIGURATION_LOAD_EVENT] = event_data async def _metrics_event(self, event_data: Dict[str, str]): async with self.lock: logger.info(f"Metrics Event: {event_data}") self.cache[METRICS_EVENT] = event_data # Custom ping handler def on_ping(self, data: bytes): logger.info(f"Received ping with data: {data}") # Send a custom readiness message, "ready", as expected by protocol self.write_message(CONNECTION_READY_MESSAGE)