master
js 425 lines 14.3 KB
Raw
1 #!/usr/bin/env node
2
3 const WebSocket = require('ws');
4 const path = require('path');
5
6 // Get program name for logs
7 const PROGRAM_NAME = path.basename(process.argv[1] || 'nd-mcp-nodejs');
8
9 function usage() {
10 console.error(`${PROGRAM_NAME}: Usage: ${PROGRAM_NAME} [--bearer TOKEN] ws://host/path`);
11 process.exit(1);
12 }
13
14 const parsedArgs = process.argv.slice(2);
15 let targetURL = '';
16 let bearerToken = '';
17
18 for (let i = 0; i < parsedArgs.length;) {
19 const arg = parsedArgs[i];
20
21 if (arg === '--bearer') {
22 if (i + 1 >= parsedArgs.length) usage();
23 bearerToken = parsedArgs[i + 1].trim();
24 i += 2;
25 }
26 else if (arg.startsWith('--bearer=')) {
27 bearerToken = arg.substring('--bearer='.length).trim();
28 i += 1;
29 }
30 else {
31 if (targetURL) usage();
32 targetURL = arg;
33 i += 1;
34 }
35 }
36
37 if (!targetURL) usage();
38
39 if (!bearerToken) {
40 const envToken = process.env.ND_MCP_BEARER_TOKEN;
41 if (envToken) bearerToken = envToken.trim();
42 }
43
44 if (bearerToken) {
45 console.error(`${PROGRAM_NAME}: Authorization header enabled for MCP connection`);
46 }
47
48 // Reconnection settings
49 const MAX_RECONNECT_DELAY_MS = 60000; // 60 seconds
50 const BASE_DELAY_MS = 1000; // 1 second
51 const CONNECTION_TIMEOUT_MS = 5000; // 5 seconds timeout for initial connection
52 let reconnectAttempt = 0;
53 let ws = null;
54 let messageQueue = [];
55 let reconnectTimeout = null;
56 let stdinActive = true;
57 let connectingInProgress = false;
58 let pendingRequests = new Map(); // Store pending requests by their IDs
59
60 // Parse JSON-RPC message and extract ID if present
61 function parseJsonRpcMessage(message) {
62 try {
63 const data = JSON.parse(message);
64 if (data && typeof data === 'object' && data.jsonrpc === '2.0') {
65 return [data.id, data.method];
66 }
67 return [null, null];
68 } catch (err) {
69 return [null, null];
70 }
71 }
72
73 // Create a JSON-RPC error response
74 function createJsonRpcError(id, code, message, data = null) {
75 const response = {
76 jsonrpc: '2.0',
77 id: id,
78 error: {
79 code: code,
80 message: message
81 }
82 };
83
84 if (data !== null) {
85 response.error.data = data;
86 }
87
88 return JSON.stringify(response);
89 }
90
91 // Handle request timeout for connection establishment only
92 function handleRequestTimeout(msgId) {
93 // This timeout ONLY applies if we're still not connected after waiting
94 if (!ws || ws.readyState !== WebSocket.OPEN) {
95 if (pendingRequests.has(msgId)) {
96 console.error(`${PROGRAM_NAME}: Connection timeout for request ID ${msgId}, sending error response`);
97
98 // Create and send error response
99 const errorResponse = createJsonRpcError(
100 msgId,
101 -32000, // Server error code
102 "MCP server connection failed",
103 { details: "Could not establish connection to Netdata within timeout period" }
104 );
105
106 try {
107 process.stdout.write(errorResponse + "\n");
108 } catch (err) {
109 console.error(`${PROGRAM_NAME}: ERROR: Failed to write error response to stdout: ${err.message}`);
110 }
111
112 // Get the original message from pendingRequests
113 const originalMessage = pendingRequests.get(msgId);
114
115 // Remove from pending requests
116 pendingRequests.delete(msgId);
117
118 // Also remove from messageQueue if it exists there
119 const msgIndex = messageQueue.indexOf(originalMessage);
120 if (msgIndex !== -1) {
121 messageQueue.splice(msgIndex, 1);
122 console.error(`${PROGRAM_NAME}: Removed timed-out request from message queue`);
123 }
124 }
125 } else {
126 // If we get here, it means the connection was established before the timeout
127 // Just clear the pending request without any action
128 if (pendingRequests.has(msgId)) {
129 // Connection succeeded in time, just remove the tracking
130 pendingRequests.delete(msgId);
131 }
132 }
133 }
134
135 // Set up stdin processing once
136 process.stdin.setEncoding('utf8');
137 process.stdin.on('data', (data) => {
138 try {
139 // Flag to check if we received actual data
140 let hasContent = false;
141
142 data.split(/\r?\n/).forEach(line => {
143 if (line.trim() !== '') {
144 hasContent = true;
145
146 // Parse the message for JSON-RPC ID
147 const [msgId, method] = parseJsonRpcMessage(line);
148
149 if (ws && ws.readyState === WebSocket.OPEN) {
150 // If connected, send immediately and track requests with IDs
151 ws.send(line);
152
153 // If this is a request with an ID, track it (but don't set a timeout)
154 if (msgId !== null && msgId !== undefined) {
155 pendingRequests.set(msgId, line);
156 }
157 } else {
158 // Queue the message if not connected
159 messageQueue.push(line);
160
161 // If this is a request with an ID, set a timeout for CONNECTION establishment
162 if (msgId !== null && msgId !== undefined) {
163 console.error(`${PROGRAM_NAME}: Received request with ID ${msgId}, setting connection timeout`);
164 pendingRequests.set(msgId, line);
165
166 // Set timeout to send error if connection not established quickly
167 // This ONLY applies to connection establishment, not request processing
168 setTimeout(() => handleRequestTimeout(msgId), CONNECTION_TIMEOUT_MS);
169 }
170 }
171 }
172 });
173
174 // If disconnected and we got content, try to reconnect immediately
175 if (hasContent && (!ws || ws.readyState !== WebSocket.OPEN)) {
176 console.error(`${PROGRAM_NAME}: Received stdin data, attempting immediate reconnection`);
177 // Clear any pending reconnect timeout
178 if (reconnectTimeout) {
179 clearTimeout(reconnectTimeout);
180 reconnectTimeout = null;
181 }
182
183 // Only attempt connection if not already connecting
184 if (!connectingInProgress) {
185 connect(true); // immediate=true
186 }
187 }
188 } catch (err) {
189 console.error(`${PROGRAM_NAME}: ERROR: Failed to process data: ${err.message}`);
190 }
191 });
192
193 process.stdin.on('error', (err) => {
194 console.error(`${PROGRAM_NAME}: ERROR: stdin: ${err.message}`);
195 stdinActive = false;
196 process.exit(1);
197 });
198
199 process.stdin.on('end', () => {
200 console.error(`${PROGRAM_NAME}: End of stdin, will exit when WebSocket disconnects`);
201 stdinActive = false;
202
203 // If WebSocket is not active, exit immediately
204 if (!ws || ws.readyState !== WebSocket.OPEN) {
205 process.exit(0);
206 }
207 });
208
209 function connect(immediate = false) {
210 // Clear any existing reconnect timeout
211 if (reconnectTimeout) {
212 clearTimeout(reconnectTimeout);
213 reconnectTimeout = null;
214 }
215
216 // If stdin is no longer active (closed/error) and we're disconnected, exit
217 if (!stdinActive && (!ws || ws.readyState !== WebSocket.OPEN)) {
218 console.error(`${PROGRAM_NAME}: Stdin closed and WebSocket disconnected, exiting`);
219 process.exit(0);
220 }
221
222 if (immediate || reconnectAttempt === 0) {
223 attemptConnection();
224 } else {
225 // Exponential backoff with jitter
226 const delay = Math.min(
227 MAX_RECONNECT_DELAY_MS,
228 BASE_DELAY_MS * Math.pow(2, reconnectAttempt - 1) * (0.5 + Math.random())
229 );
230
231 console.error(`${PROGRAM_NAME}: Reconnecting in ${(delay/1000).toFixed(1)} seconds (attempt ${reconnectAttempt+1})...`);
232 reconnectTimeout = setTimeout(() => {
233 reconnectTimeout = null;
234 attemptConnection();
235 }, delay);
236 }
237 }
238
239 function attemptConnection() {
240 // Prevent multiple concurrent connection attempts
241 if (connectingInProgress) {
242 return;
243 }
244
245 connectingInProgress = true;
246 console.error(`${PROGRAM_NAME}: Connecting to ${targetURL}...`);
247
248 // Close any existing websocket
249 if (ws) {
250 try {
251 ws.terminate();
252 } catch (e) {
253 // Ignore errors
254 }
255 ws = null;
256 }
257
258 const wsOptions = {
259 perMessageDeflate: true,
260 maxPayload: 16 * 1024 * 1024, // 16MB to match Netdata's limits
261 handshakeTimeout: 10000, // 10 seconds for initial handshake
262 followRedirects: true, // Follow HTTP redirects
263 // Keep the connection alive with pings
264 pingInterval: 30000, // 30 seconds
265 pingTimeout: 10000 // 10 seconds to wait for pong
266 };
267
268 if (bearerToken) {
269 wsOptions.headers = {
270 Authorization: `Bearer ${bearerToken}`
271 };
272 }
273
274 ws = new WebSocket(targetURL, wsOptions);
275
276 // Set a timeout for initial connection
277 const connectionTimeout = setTimeout(() => {
278 if (ws && ws.readyState !== WebSocket.OPEN) {
279 console.error(`${PROGRAM_NAME}: Connection timeout, closing...`);
280 ws.terminate();
281 }
282 }, 15000); // 15 seconds connection timeout
283
284 ws.on('open', () => {
285 clearTimeout(connectionTimeout);
286 console.error(`${PROGRAM_NAME}: Connected`);
287 connectingInProgress = false;
288 reconnectAttempt = 0; // Reset counter on successful connection
289
290 // Process any queued messages
291 if (messageQueue.length > 0) {
292 console.error(`${PROGRAM_NAME}: Sending ${messageQueue.length} queued message(s)`);
293 const queueCopy = [...messageQueue];
294 messageQueue = [];
295 queueCopy.forEach(msg => {
296 try {
297 ws.send(msg);
298
299 // If this was a pending request with an ID, remove it from pending
300 const [msgId] = parseJsonRpcMessage(msg);
301 if (msgId !== null && pendingRequests.has(msgId)) {
302 pendingRequests.delete(msgId);
303 }
304
305 // Ensure we don't let the message queue grow indefinitely
306 if (messageQueue.length > 1000) {
307 console.error(`${PROGRAM_NAME}: Message queue too large (${messageQueue.length}), trimming older messages`);
308 messageQueue = messageQueue.slice(-500); // Keep only the 500 most recent messages
309 }
310 } catch (err) {
311 console.error(`${PROGRAM_NAME}: ERROR: Failed to send queued message: ${err.message}`);
312 // Re-queue the message if connection is still open
313 if (ws.readyState === WebSocket.OPEN) {
314 messageQueue.push(msg);
315 }
316 }
317 });
318 }
319 });
320
321 ws.on('message', (message) => {
322 try {
323 process.stdout.write(message + "\n");
324
325 // If this is a response to a request, check if we can remove it from pending
326 try {
327 const data = JSON.parse(message);
328 if (data && typeof data === 'object' && data.jsonrpc === '2.0' && data.id !== undefined) {
329 if (pendingRequests.has(data.id)) {
330 pendingRequests.delete(data.id);
331 }
332 }
333 } catch (err) {
334 // Ignore parsing errors
335 }
336
337 // Memory management: Don't let pendingRequests grow indefinitely
338 if (pendingRequests.size > 1000) {
339 console.error(`${PROGRAM_NAME}: Too many pending requests (${pendingRequests.size}), clearing older ones`);
340 // Since Map iteration is in insertion order, we keep the newest entries
341 const entries = Array.from(pendingRequests.entries());
342 pendingRequests.clear();
343 entries.slice(-500).forEach(([id, msg]) => pendingRequests.set(id, msg));
344 }
345 } catch (err) {
346 console.error(`${PROGRAM_NAME}: ERROR: Failed to write to stdout: ${err.message}`);
347 // If stdout is broken, exit
348 process.exit(1);
349 }
350 });
351
352 ws.on('close', (code, reason) => {
353 clearTimeout(connectionTimeout);
354 console.error(`${PROGRAM_NAME}: WebSocket closed with code ${code}${reason ? ': ' + reason : ''}`);
355 connectingInProgress = false;
356
357 // Only increment reconnect attempt if we're still supposed to be running
358 if (stdinActive) {
359 reconnectAttempt++;
360 connect();
361 } else {
362 console.error(`${PROGRAM_NAME}: Stdin closed, exiting`);
363 process.exit(0);
364 }
365 });
366
367 ws.on('error', (err) => {
368 console.error(`${PROGRAM_NAME}: ERROR: WebSocket error: ${err.message}`);
369 connectingInProgress = false;
370 // The close event should handle reconnection
371 });
372
373 // Verify the connection is healthy periodically
374 ws.on('pong', () => {
375 // Connection is alive
376 });
377 }
378
379 // Handle process termination
380 process.on('SIGINT', () => {
381 console.error(`${PROGRAM_NAME}: Received SIGINT, shutting down`);
382 // Clear any pending reconnect
383 if (reconnectTimeout) {
384 clearTimeout(reconnectTimeout);
385 reconnectTimeout = null;
386 }
387
388 if (ws) {
389 try {
390 ws.close(1000, 'Normal closure');
391 } catch (e) {
392 // Ignore errors during cleanup
393 try {
394 ws.terminate();
395 } catch (e) {
396 // Last resort
397 }
398 }
399 }
400 process.exit(0);
401 });
402
403 process.on('SIGTERM', () => {
404 console.error(`${PROGRAM_NAME}: Received SIGTERM, shutting down`);
405 // Same cleanup as SIGINT
406 if (reconnectTimeout) {
407 clearTimeout(reconnectTimeout);
408 }
409
410 if (ws) {
411 try {
412 ws.close(1000, 'Normal closure');
413 } catch (e) {
414 try {
415 ws.terminate();
416 } catch (e) {
417 // Ignore
418 }
419 }
420 }
421 process.exit(0);
422 });
423
424 // Start the connection process
425 connect();