UNPKG

vibe-coder-mcp

Version:

Production-ready MCP server with complete agent integration, multi-transport support, and comprehensive development automation tools for AI-assisted workflows.

199 lines (198 loc) 8.02 kB
import logger from '../../logger.js'; import { jobManager } from '../job-manager/index.js'; import { createJobStatusMessage } from '../job-manager/jobStatusMessage.js'; class SseNotifier { connections = new Map(); registerConnection(sessionId, res) { if (this.connections.has(sessionId)) { logger.warn({ sessionId }, `SSE connection already registered for this session. Overwriting.`); const oldRes = this.connections.get(sessionId); try { oldRes?.end(); } catch (e) { logger.error({ err: e, sessionId }, `Error closing previous SSE connection.`); } } res.writeHead(200, { 'Content-Type': 'text/event-stream', 'Cache-Control': 'no-cache', 'Connection': 'keep-alive', 'Access-Control-Allow-Origin': '*', }); res.flushHeaders(); this.connections.set(sessionId, res); logger.info({ sessionId }, `Registered new SSE connection.`); this.sendMessage(sessionId, 'connected', { message: 'SSE connection established.' }); res.on('close', () => { this.unregisterConnection(sessionId); logger.info({ sessionId }, `SSE connection closed by client.`); }); const keepAliveInterval = setInterval(() => { if (!this.connections.has(sessionId)) { clearInterval(keepAliveInterval); return; } const currentRes = this.connections.get(sessionId); if (currentRes && !currentRes.writableEnded) { try { currentRes.write(': keep-alive\n\n'); } catch (e) { logger.warn({ err: e, sessionId }, 'Error writing keep-alive. Clearing interval and unregistering.'); clearInterval(keepAliveInterval); this.unregisterConnection(sessionId); } } else { clearInterval(keepAliveInterval); if (this.connections.has(sessionId)) { this.unregisterConnection(sessionId); } } }, 20000); } unregisterConnection(sessionId) { const res = this.connections.get(sessionId); if (res) { try { if (!res.writableEnded) { res.end(); } } catch (e) { logger.error({ err: e, sessionId }, `Error ending SSE connection during unregister.`); } this.connections.delete(sessionId); logger.info({ sessionId }, `Unregistered SSE connection.`); } else { logger.warn({ sessionId }, `Attempted to unregister non-existent SSE connection.`); } } sendMessage(sessionId, event, data) { const res = this.connections.get(sessionId); if (res && !res.writableEnded) { try { res.write(`event: ${event}\n`); res.write(`data: ${JSON.stringify(data)}\n\n`); logger.debug({ sessionId, event, data }, `Sent SSE message.`); } catch (e) { logger.error({ err: e, sessionId, event }, `Failed to send SSE message.`); this.unregisterConnection(sessionId); } } else if (!res) { logger.warn({ sessionId, event }, `Attempted to send SSE message to non-existent or closed connection.`); } } sendProgress(sessionId, jobId, status, message, progress) { if (!sessionId) { logger.warn({ jobId, status, message }, "Cannot send progress update: Missing sessionId."); return; } const job = jobManager.getJob(jobId, false); if (!job) { logger.warn({ jobId, status, message }, "Cannot send progress update: Job not found."); return; } const statusMessage = createJobStatusMessage(jobId, job.toolName, status, message, progress, job.createdAt, job.updatedAt); if (sessionId === 'stdio-session' || sessionId === 'placeholder-session-id' || sessionId === 'unknown-session') { const updateStartTime = Date.now(); try { jobManager.updateJobStatus(jobId, status, message); if (message && progress !== undefined) { const progressOutput = { type: 'progress', jobId, tool: job.toolName, status, message, progress, timestamp: new Date().toISOString() }; process.stderr.write(`[PROGRESS] ${JSON.stringify(progressOutput)}\n`); } const updateTime = Date.now() - updateStartTime; logger.info({ jobId, status, message, sessionId, updateTime, performance: { jobUpdateLatency: updateTime, timestamp: new Date().toISOString(), statusType: status } }, "Job status update: successful stdio session update with stderr emission"); return; } catch (error) { const updateTime = Date.now() - updateStartTime; logger.error({ err: error, jobId, status, message, sessionId, updateTime, errorType: error instanceof Error ? error.constructor.name : 'Unknown', errorMessage: error instanceof Error ? error.message : String(error) }, "Job status update: failed to update job status for stdio session"); return; } } if (this.connections.has(sessionId)) { this.sendMessage(sessionId, 'jobProgress', statusMessage); } else { logger.warn({ jobId, status, message, sessionId }, "Cannot send SSE progress: No active connection for session."); } } sendJobResult(sessionId, jobId, result) { if (!sessionId) { logger.warn({ jobId }, "Cannot send job result: Missing sessionId."); return; } if (sessionId === 'stdio-session' || sessionId === 'placeholder-session-id' || sessionId === 'unknown-session') { logger.debug({ jobId, sessionId }, "Skipping SSE job result for stdio session"); return; } if (this.connections.has(sessionId)) { this.sendMessage(sessionId, 'result', { jobId, result }); } else { logger.warn({ jobId, sessionId }, "Cannot send SSE job result: No active connection for session."); } } async sendEvent(sessionId, event, data) { this.sendMessage(sessionId, event, data); } async broadcastEvent(event, data) { for (const sessionId of this.connections.keys()) { this.sendMessage(sessionId, event, data); } } getConnectionCount() { return this.connections.size; } closeAllConnections() { logger.info(`Closing all ${this.connections.size} active SSE connections...`); this.connections.forEach((res, sessionId) => { this.unregisterConnection(sessionId); }); this.connections.clear(); logger.info(`All SSE connections closed.`); } } export const sseNotifier = new SseNotifier(); process.on('SIGINT', () => { sseNotifier.closeAllConnections(); process.exit(0); }); process.on('SIGTERM', () => { sseNotifier.closeAllConnections(); process.exit(0); });