| 1 | #!/usr/bin/env python3 |
| 2 | # -*- coding: utf-8 -*- |
| 3 | |
| 4 | import sys |
| 5 | import asyncio |
| 6 | import websockets |
| 7 | import os |
| 8 | import random |
| 9 | import time |
| 10 | import signal |
| 11 | import json |
| 12 | |
| 13 | # Get program name for logs |
| 14 | PROGRAM_NAME = os.path.basename(sys.argv[0]) if len(sys.argv) > 0 else "nd-mcp-python" |
| 15 | |
| 16 | # Global flag to track if we're exiting due to a stdin issue |
| 17 | STDIN_ERROR = False |
| 18 | |
| 19 | # Timeout for connection attempts when a message with ID is waiting |
| 20 | CONNECTION_TIMEOUT = 5 # seconds |
| 21 | |
| 22 | # Parse JSON-RPC message and extract ID if present |
| 23 | def parse_jsonrpc_message(message): |
| 24 | try: |
| 25 | data = json.loads(message) |
| 26 | if isinstance(data, dict) and "jsonrpc" in data and data.get("jsonrpc") == "2.0": |
| 27 | return data.get("id"), data.get("method") |
| 28 | return None, None |
| 29 | except json.JSONDecodeError: |
| 30 | return None, None |
| 31 | |
| 32 | # Create a JSON-RPC error response |
| 33 | def create_jsonrpc_error(id, code, message, data=None): |
| 34 | response = { |
| 35 | "jsonrpc": "2.0", |
| 36 | "id": id, |
| 37 | "error": { |
| 38 | "code": code, |
| 39 | "message": message |
| 40 | } |
| 41 | } |
| 42 | if data is not None: |
| 43 | response["error"]["data"] = data |
| 44 | return json.dumps(response) |
| 45 | |
| 46 | async def connect_with_backoff(uri, bearer_token): |
| 47 | max_delay = 60 # Maximum delay between reconnections in seconds |
| 48 | base_delay = 1 # Initial delay in seconds |
| 49 | retry_count = 0 |
| 50 | |
| 51 | # Message queue for storing messages during disconnections |
| 52 | stdin_queue = asyncio.Queue() |
| 53 | |
| 54 | # Dictionary to track pending requests with their IDs and timers |
| 55 | pending_requests = {} |
| 56 | |
| 57 | # Set up stdin reader once |
| 58 | async def read_stdin(): |
| 59 | global STDIN_ERROR |
| 60 | try: |
| 61 | loop = asyncio.get_running_loop() |
| 62 | reader = asyncio.StreamReader() |
| 63 | await loop.connect_read_pipe(lambda: asyncio.StreamReaderProtocol(reader), sys.stdin) |
| 64 | |
| 65 | while True: |
| 66 | line = await reader.readline() |
| 67 | if not line: |
| 68 | print(f"{PROGRAM_NAME}: End of stdin, exiting...", file=sys.stderr) |
| 69 | STDIN_ERROR = True |
| 70 | # Signal the main event loop to exit |
| 71 | loop.stop() |
| 72 | return |
| 73 | |
| 74 | # Process the received line |
| 75 | line_text = line.decode().strip() |
| 76 | if not line_text: |
| 77 | continue |
| 78 | |
| 79 | # Parse JSON-RPC message |
| 80 | msg_id, msg_method = parse_jsonrpc_message(line_text) |
| 81 | |
| 82 | # Store message in queue |
| 83 | await stdin_queue.put((line_text, msg_id, msg_method)) |
| 84 | |
| 85 | # If we're disconnected, check if we need to respond with error |
| 86 | if retry_count > 0: |
| 87 | print(f"{PROGRAM_NAME}: Received stdin data, attempting immediate reconnection", file=sys.stderr) |
| 88 | retry_event.set() |
| 89 | |
| 90 | if msg_id is not None: |
| 91 | # Set timer for this request |
| 92 | pending_requests[msg_id] = { |
| 93 | "message": line_text, |
| 94 | "timer": asyncio.create_task(handle_request_timeout(msg_id, CONNECTION_TIMEOUT)) |
| 95 | } |
| 96 | |
| 97 | except Exception as e: |
| 98 | print(f"{PROGRAM_NAME}: ERROR: stdin reader: {e}", file=sys.stderr) |
| 99 | STDIN_ERROR = True |
| 100 | # Signal the main event loop to exit |
| 101 | loop.stop() |
| 102 | return |
| 103 | |
| 104 | # Handler for request timeout - ONLY for connection establishment |
| 105 | async def handle_request_timeout(msg_id, timeout): |
| 106 | await asyncio.sleep(timeout) |
| 107 | |
| 108 | # If we're still disconnected and the request is still pending |
| 109 | if retry_count > 0 and msg_id in pending_requests: |
| 110 | print(f"{PROGRAM_NAME}: Connection timeout for request ID {msg_id}, sending error response", file=sys.stderr) |
| 111 | |
| 112 | # Create and send error response |
| 113 | error_response = create_jsonrpc_error( |
| 114 | msg_id, |
| 115 | -32000, # Server error code |
| 116 | "MCP server connection failed", |
| 117 | {"details": "Could not establish connection to Netdata within timeout period"} |
| 118 | ) |
| 119 | print(error_response, flush=True) |
| 120 | |
| 121 | # Mark this message as timed out |
| 122 | if msg_id in pending_requests: |
| 123 | pending_requests[msg_id]["timed_out"] = True |
| 124 | |
| 125 | # We'll keep it in pending_requests so we don't resend it, but mark it as handled |
| 126 | print(f"{PROGRAM_NAME}: Marked request ID {msg_id} as timed out - it will not be sent to server", file=sys.stderr) |
| 127 | elif msg_id in pending_requests: |
| 128 | # Connection was established before timeout, just clean up the timer |
| 129 | print(f"{PROGRAM_NAME}: Connection established before timeout for request ID {msg_id}", file=sys.stderr) |
| 130 | |
| 131 | # Create an event for signaling reconnection |
| 132 | retry_event = asyncio.Event() |
| 133 | |
| 134 | # Start reading stdin in the background |
| 135 | stdin_task = asyncio.create_task(read_stdin()) |
| 136 | |
| 137 | while True: |
| 138 | if STDIN_ERROR: |
| 139 | print(f"{PROGRAM_NAME}: Stdin error detected, exiting", file=sys.stderr) |
| 140 | return |
| 141 | |
| 142 | try: |
| 143 | # Calculate backoff delay with jitter |
| 144 | delay = min(max_delay, base_delay * (2 ** retry_count) * (0.5 + random.random())) |
| 145 | |
| 146 | if retry_count > 0: |
| 147 | print(f"{PROGRAM_NAME}: Reconnecting in {delay:.1f} seconds (attempt {retry_count+1})...", file=sys.stderr) |
| 148 | |
| 149 | # Create a wait task, but also break on retry_event being set |
| 150 | try: |
| 151 | # Wait for the delay or until retry_event is set |
| 152 | wait_task = asyncio.create_task(asyncio.sleep(delay)) |
| 153 | retry_task = asyncio.create_task(retry_event.wait()) |
| 154 | |
| 155 | done, pending = await asyncio.wait( |
| 156 | [wait_task, retry_task], |
| 157 | return_when=asyncio.FIRST_COMPLETED |
| 158 | ) |
| 159 | |
| 160 | # Cancel the pending task |
| 161 | for task in pending: |
| 162 | task.cancel() |
| 163 | |
| 164 | # Clear the event |
| 165 | retry_event.clear() |
| 166 | |
| 167 | except asyncio.CancelledError: |
| 168 | # Cancel the in-flight wait/retry tasks so they don't |
| 169 | # outlive this coroutine, then propagate the cancellation |
| 170 | # so the parent reconnect loop can exit cleanly. |
| 171 | wait_task.cancel() |
| 172 | retry_task.cancel() |
| 173 | retry_event.clear() |
| 174 | raise |
| 175 | |
| 176 | print(f"{PROGRAM_NAME}: Connecting to {uri}...", file=sys.stderr) |
| 177 | |
| 178 | try: |
| 179 | # Connect with timeout |
| 180 | # In newer versions of websockets, connect() is already awaitable |
| 181 | connect_kwargs = { |
| 182 | "compression": 'deflate', |
| 183 | "max_size": 16*1024*1024, |
| 184 | "ping_interval": 30, |
| 185 | "ping_timeout": 10, |
| 186 | "close_timeout": 5 |
| 187 | } |
| 188 | |
| 189 | if bearer_token: |
| 190 | connect_kwargs["extra_headers"] = { |
| 191 | "Authorization": f"Bearer {bearer_token}" |
| 192 | } |
| 193 | |
| 194 | ws = await asyncio.wait_for( |
| 195 | websockets.connect( |
| 196 | uri, |
| 197 | **connect_kwargs |
| 198 | ), |
| 199 | timeout=15 # 15 second timeout |
| 200 | ) |
| 201 | except asyncio.TimeoutError: |
| 202 | raise Exception("Connection timeout") |
| 203 | |
| 204 | print(f"{PROGRAM_NAME}: Connected", file=sys.stderr) |
| 205 | retry_count = 0 # Reset retry counter on successful connection |
| 206 | |
| 207 | # Clear all pending request timers |
| 208 | for req_id, req_data in list(pending_requests.items()): |
| 209 | if "timer" in req_data and not req_data["timer"].done(): |
| 210 | req_data["timer"].cancel() |
| 211 | |
| 212 | # We don't clear pending_requests entirely here, we just clear the timers |
| 213 | # and keep tracking the requests until we get responses |
| 214 | |
| 215 | # Memory management: Limit the size of pending_requests to avoid memory leaks |
| 216 | if len(pending_requests) > 1000: |
| 217 | print(f"{PROGRAM_NAME}: Too many pending requests ({len(pending_requests)}), cleaning up", file=sys.stderr) |
| 218 | pending_requests.clear() # In extreme case, clear everything |
| 219 | |
| 220 | # Processor for stdin messages |
| 221 | async def process_stdin(): |
| 222 | try: |
| 223 | while True: |
| 224 | # Get message from queue |
| 225 | line_data = await stdin_queue.get() |
| 226 | line, msg_id, _ = line_data |
| 227 | |
| 228 | # Skip messages that have already timed out and received error responses |
| 229 | if msg_id is not None and msg_id in pending_requests and pending_requests[msg_id].get("timed_out", False): |
| 230 | print(f"{PROGRAM_NAME}: Skipping previously timed-out request with ID {msg_id}", file=sys.stderr) |
| 231 | stdin_queue.task_done() |
| 232 | continue |
| 233 | |
| 234 | # If this is a request with an ID, track it (but no timeout since we're connected) |
| 235 | if msg_id is not None and msg_id not in pending_requests: |
| 236 | pending_requests[msg_id] = { |
| 237 | "message": line, |
| 238 | "sent_time": time.time() |
| 239 | } |
| 240 | |
| 241 | try: |
| 242 | # Send to WebSocket |
| 243 | await ws.send(line) |
| 244 | stdin_queue.task_done() |
| 245 | except websockets.exceptions.ConnectionClosed: |
| 246 | # Put the message back in the queue, unless it already timed out |
| 247 | if not (msg_id is not None and msg_id in pending_requests and pending_requests[msg_id].get("timed_out", False)): |
| 248 | await stdin_queue.put(line_data) |
| 249 | # Re-raise to trigger reconnection |
| 250 | raise |
| 251 | except Exception as e: |
| 252 | # In newer websockets versions, we need to check for close attribute differently |
| 253 | try: |
| 254 | # Check if the connection is still open |
| 255 | if hasattr(ws, 'closed') and not ws.closed: |
| 256 | print(f"{PROGRAM_NAME}: ERROR: stdin processor: {e}", file=sys.stderr) |
| 257 | elif hasattr(ws, 'protocol') and not ws.protocol.closed: |
| 258 | print(f"{PROGRAM_NAME}: ERROR: stdin processor: {e}", file=sys.stderr) |
| 259 | else: |
| 260 | print(f"{PROGRAM_NAME}: Websocket connection closed: {e}", file=sys.stderr) |
| 261 | except: |
| 262 | # If we can't check closed state, just log the error |
| 263 | print(f"{PROGRAM_NAME}: ERROR: stdin processor: {e}", file=sys.stderr) |
| 264 | # Don't propagate the exception, let the connection close |
| 265 | # and reconnection will be triggered |
| 266 | |
| 267 | # Processor for WebSocket messages |
| 268 | async def process_websocket(): |
| 269 | try: |
| 270 | async for message in ws: |
| 271 | # Forward WebSocket messages to stdout |
| 272 | print(message, flush=True) |
| 273 | |
| 274 | # Check if this is a response to a tracked request |
| 275 | try: |
| 276 | data = json.loads(message) |
| 277 | if isinstance(data, dict) and "jsonrpc" in data and data.get("jsonrpc") == "2.0" and "id" in data: |
| 278 | msg_id = data.get("id") |
| 279 | if msg_id in pending_requests: |
| 280 | del pending_requests[msg_id] |
| 281 | except: |
| 282 | pass |
| 283 | |
| 284 | # Memory management: Limit the size of stdin_queue to avoid memory leaks |
| 285 | if stdin_queue.qsize() > 1000: |
| 286 | print(f"{PROGRAM_NAME}: WARNING: Very large stdin queue ({stdin_queue.qsize()}), this may indicate a problem", file=sys.stderr) |
| 287 | except Exception as e: |
| 288 | # In newer websockets versions, we need to check for close attribute differently |
| 289 | try: |
| 290 | # Check if the connection is still open |
| 291 | if hasattr(ws, 'closed') and not ws.closed: |
| 292 | print(f"{PROGRAM_NAME}: ERROR: websocket processor: {e}", file=sys.stderr) |
| 293 | elif hasattr(ws, 'protocol') and not ws.protocol.closed: |
| 294 | print(f"{PROGRAM_NAME}: ERROR: websocket processor: {e}", file=sys.stderr) |
| 295 | else: |
| 296 | print(f"{PROGRAM_NAME}: Websocket connection closed: {e}", file=sys.stderr) |
| 297 | except: |
| 298 | # If we can't check closed state, just log the error |
| 299 | print(f"{PROGRAM_NAME}: ERROR: websocket processor: {e}", file=sys.stderr) |
| 300 | # Don't propagate the exception, let the connection close |
| 301 | # and reconnection will be triggered |
| 302 | |
| 303 | # Run both tasks concurrently |
| 304 | stdin_processor = asyncio.create_task(process_stdin()) |
| 305 | ws_processor = asyncio.create_task(process_websocket()) |
| 306 | |
| 307 | # Wait for either task to complete (which means a failure) |
| 308 | done, pending = await asyncio.wait( |
| 309 | [stdin_processor, ws_processor], |
| 310 | return_when=asyncio.FIRST_COMPLETED |
| 311 | ) |
| 312 | |
| 313 | if pending: |
| 314 | # Drain cancelled children without suppressing cancellation of this coroutine. |
| 315 | for task in pending: |
| 316 | task.cancel() |
| 317 | results = await asyncio.gather(*pending, return_exceptions=True) |
| 318 | for result in results: |
| 319 | if isinstance(result, asyncio.CancelledError): |
| 320 | continue |
| 321 | if isinstance(result, BaseException): |
| 322 | raise result |
| 323 | |
| 324 | # Ensure WebSocket is closed |
| 325 | try: |
| 326 | # Check if websocket is already closed |
| 327 | is_closed = False |
| 328 | if hasattr(ws, 'closed'): |
| 329 | is_closed = ws.closed |
| 330 | elif hasattr(ws, 'protocol'): |
| 331 | is_closed = ws.protocol.closed |
| 332 | |
| 333 | if not is_closed: |
| 334 | await ws.close() |
| 335 | except Exception as e: |
| 336 | print(f"{PROGRAM_NAME}: Error closing websocket: {e}", file=sys.stderr) |
| 337 | |
| 338 | except (websockets.exceptions.ConnectionClosed, websockets.exceptions.WebSocketException) as e: |
| 339 | print(f"{PROGRAM_NAME}: WebSocket error: {e}", file=sys.stderr) |
| 340 | retry_count += 1 |
| 341 | except Exception as e: |
| 342 | print(f"{PROGRAM_NAME}: Unexpected error: {e}", file=sys.stderr) |
| 343 | retry_count += 1 |
| 344 | |
| 345 | def usage(): |
| 346 | print(f"{PROGRAM_NAME}: Usage: {PROGRAM_NAME} [--bearer TOKEN] ws://host/path", file=sys.stderr) |
| 347 | sys.exit(1) |
| 348 | |
| 349 | |
| 350 | def parse_args(argv): |
| 351 | target = None |
| 352 | bearer = None |
| 353 | idx = 0 |
| 354 | |
| 355 | while idx < len(argv): |
| 356 | arg = argv[idx] |
| 357 | if arg == '--bearer': |
| 358 | if idx + 1 >= len(argv): |
| 359 | usage() |
| 360 | bearer = argv[idx + 1].strip() |
| 361 | idx += 2 |
| 362 | elif arg.startswith('--bearer='): |
| 363 | bearer = arg.split('=', 1)[1].strip() |
| 364 | idx += 1 |
| 365 | else: |
| 366 | if target is not None: |
| 367 | usage() |
| 368 | target = arg |
| 369 | idx += 1 |
| 370 | |
| 371 | if not target: |
| 372 | usage() |
| 373 | |
| 374 | return target, bearer |
| 375 | |
| 376 | |
| 377 | def main(): |
| 378 | target_uri, bearer_token = parse_args(sys.argv[1:]) |
| 379 | |
| 380 | if not bearer_token: |
| 381 | env_token = os.environ.get("ND_MCP_BEARER_TOKEN", "") |
| 382 | if env_token: |
| 383 | bearer_token = env_token.strip() |
| 384 | |
| 385 | if bearer_token: |
| 386 | print(f"{PROGRAM_NAME}: Authorization header enabled for MCP connection", file=sys.stderr) |
| 387 | |
| 388 | # Set up signal handling |
| 389 | def signal_handler(sig, frame): |
| 390 | print(f"{PROGRAM_NAME}: Received signal {sig}, exiting", file=sys.stderr) |
| 391 | sys.exit(0) |
| 392 | |
| 393 | signal.signal(signal.SIGINT, signal_handler) |
| 394 | signal.signal(signal.SIGTERM, signal_handler) |
| 395 | |
| 396 | try: |
| 397 | asyncio.run(connect_with_backoff(target_uri, bearer_token)) |
| 398 | except KeyboardInterrupt: |
| 399 | print(f"{PROGRAM_NAME}: Interrupted by user, exiting", file=sys.stderr) |
| 400 | |
| 401 | if STDIN_ERROR: |
| 402 | print(f"{PROGRAM_NAME}: Exiting due to stdin error", file=sys.stderr) |
| 403 | |
| 404 | if __name__ == "__main__": |
| 405 | main() |