Implementation of MCP for LLM Observability capture to PostHig

src/server/mcpServer.ts

@@ -0,0 +1,90 @@
+import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
+import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js';
+import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
+
+import { Logger } from '../utils/logger.util';
+import { config } from '../utils/config.util';
+import { PACKAGE_NAME, VERSION } from '../utils/constants.util';
+import posthogLlmResources from '../resources/posthog-llm.resource.js';
+import posthogLlmTools from '../tools/posthog-llm.tool.js';
+
+export function createServer() {
+	const serverLogger = Logger.forContext('utils/server.util.ts', 'getServer');
+
+	// Load configuration
+	serverLogger.info('Starting MCP server initialization...');
+	config.load();
+
+	if (config.getBoolean('DEBUG')) {
+		serverLogger.debug('Debug mode enabled');
+	}
+
+	serverLogger.info(`Initializing LLM Log MCP server v${VERSION}`);
+	const server = new McpServer({
+		name: PACKAGE_NAME,
+		version: VERSION,
+	});
+
+	// Register tools and resources
+	serverLogger.info('Registering MCP tools and resources...');
+	posthogLlmTools.registerTools(server);
+	posthogLlmResources.registerResources(server);
+	serverLogger.debug('All tools and resources registered');
+
+	return server;
+}
+
+/**
+ * Graceful shutdown handler
+ */
+export function shutdownServer(
+	transports:
+		| Record<string, SSEServerTransport | StreamableHTTPServerTransport>
+		| undefined,
+) {
+	const shutdownLogger = Logger.forContext(
+		'utils/server.util.ts',
+		'shutdown',
+	);
+
+	const shutdown = async () => {
+		try {
+			shutdownLogger.info('Shutting down server...');
+
+			// Close all active transports to properly clean up resources
+			if (transports) {
+				for (const sessionId in transports) {
+					try {
+						shutdownLogger.debug(
+							`Closing transport for session ${sessionId}`,
+						);
+						const transport = transports[sessionId];
+						if (
+							transport &&
+							'close' in transport &&
+							typeof transport.close === 'function'
+						) {
+							await transport.close();
+							delete transports[sessionId];
+						}
+					} catch (error) {
+						shutdownLogger.error(
+							`Error closing transport for session ${sessionId}:`,
+							error,
+						);
+					}
+				}
+			}
+
+			shutdownLogger.info('Server shutdown complete');
+			process.exit(0);
+		} catch (err) {
+			shutdownLogger.error('Error during shutdown', err);
+			process.exit(1);
+		}
+	};
+
+	['SIGINT', 'SIGTERM'].forEach((signal) => {
+		process.on(signal as NodeJS.Signals, shutdown);
+	});
+}