Umbrella MCP Server

import { Server } from '@modelcontextprotocol/sdk/server/index.js'; import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'; import { CallToolRequestSchema, ListToolsRequestSchema, Tool, } from '@modelcontextprotocol/sdk/types.js'; import { UserSessionManager, UserCredentials } from './user-session-manager.js'; import { AuthCredentials, QueryParams, QueryParamsSchema, UMBRELLA_ENDPOINTS } from './types.js'; import * as fs from 'fs'; import * as path from 'path'; // Simple logging without JSON complexity class SimpleLogger { private logFile: string; constructor() { this.logFile = path.join(process.cwd(), 'mcp-simple.log'); this.log('=== MCP Server Starting ==='); } private log(message: string) { const timestamp = new Date().toISOString(); const logLine = `[${timestamp}] ${message}\n`; try { fs.appendFileSync(this.logFile, logLine); } catch (error) { // Ignore file write errors to prevent cascading issues } // Write to stderr for Claude Desktop console.error(`[MCP] ${message}`); } info(message: string) { this.log(`INFO: ${message}`); } error(message: string) { this.log(`ERROR: ${message}`); } auth(message: string) { this.log(`AUTH: ${message}`); } } export class RobustUmbrellaMcpServer { private server: Server; private sessionManager: UserSessionManager; private baseURL: string; private logger: SimpleLogger; constructor(baseURL: string) { this.baseURL = baseURL; this.logger = new SimpleLogger(); this.sessionManager = new UserSessionManager(baseURL); this.logger.info(`Initializing Robust Umbrella MCP Server with baseURL: ${this.baseURL}`); try { this.server = new Server( { name: 'Umbrella MCP Robust', version: '1.0.0', }, { capabilities: { tools: {}, }, } ); this.setupToolHandlers(); this.setupErrorHandling(); this.logger.info('Server initialization completed successfully'); } catch (error: any) { this.logger.error(`Server initialization failed: ${error.message}`); throw error; } } private setupErrorHandling(): void { process.on('uncaughtException', (error) => { this.logger.error(`Uncaught Exception: ${error.message}`); // Don't exit immediately, let MCP handle it }); process.on('unhandledRejection', (reason, promise) => { this.logger.error(`Unhandled Rejection: ${reason}`); // Don't exit immediately, let MCP handle it }); } private setupToolHandlers(): void { this.logger.info('Setting up tool handlers'); try { // List available tools with error handling this.server.setRequestHandler(ListToolsRequestSchema, async (request) => { this.logger.info('ListTools request received'); try { const tools: Tool[] = [ { name: 'authenticate_user', description: 'Authenticate with Umbrella Cost API using username and password', inputSchema: { type: 'object', properties: { username: { type: 'string', description: 'Your Umbrella Cost username (email address)', }, password: { type: 'string', description: 'Your Umbrella Cost password', }, }, required: ['username', 'password'], }, }, { name: 'session_status', description: 'Check authentication status', inputSchema: { type: 'object', properties: {}, required: [], }, }, { name: 'help', description: 'Get help information', inputSchema: { type: 'object', properties: {}, required: [], }, } ]; // Add a few key API tools const keyApiTools: Tool[] = [ { name: 'api_user_management_accounts', description: 'Get customer accounts information (MSP view)', inputSchema: { type: 'object', properties: { username: { type: 'string', description: 'Username for multi-tenant access (optional)', }, }, required: [], }, }, { name: 'api_recommendations_report', description: 'Get cost optimization recommendations', inputSchema: { type: 'object', properties: { username: { type: 'string', description: 'Username for multi-tenant access (optional)', }, }, required: [], }, }, { name: 'api_invoices_service_names_distinct', description: 'Get distinct service names', inputSchema: { type: 'object', properties: { limit: { type: 'string', description: 'Limit number of results', }, username: { type: 'string', description: 'Username for multi-tenant access (optional)', }, }, required: [], }, } ]; const allTools = [...tools, ...keyApiTools]; this.logger.info(`Returning ${allTools.length} tools`); return { tools: allTools, }; } catch (error: any) { this.logger.error(`Error in listTools: ${error.message}`); throw error; } }); // Handle tool calls with comprehensive error handling this.server.setRequestHandler(CallToolRequestSchema, async (request) => { const { name, arguments: args } = request.params; this.logger.info(`CallTool request: ${name}`); try { switch (name) { case 'authenticate_user': return await this.handleAuthenticateUser(args as any); case 'session_status': return this.handleSessionStatus(); case 'help': return this.handleHelp(); case 'api_user_management_accounts': return await this.handleApiCall('/user-management/accounts', args); case 'api_recommendations_report': return await this.handleApiCall('/recommendations/report', args); case 'api_invoices_service_names_distinct': return await this.handleApiCall('/invoices/service-names/distinct', args); default: this.logger.error(`Unknown tool: ${name}`); return { content: [ { type: 'text', text: `Error: Unknown tool "${name}". Available tools: authenticate_user, session_status, help, api_user_management_accounts, api_recommendations_report, api_invoices_service_names_distinct`, }, ], }; } } catch (error: any) { this.logger.error(`Error in callTool ${name}: ${error.message}`); return { content: [ { type: 'text', text: `Error executing ${name}: ${error.message}`, }, ], }; } }); this.logger.info('Tool handlers setup completed'); } catch (error: any) { this.logger.error(`Error setting up tool handlers: ${error.message}`); throw error; } } private async handleAuthenticateUser(args: { username: string; password: string }) { this.logger.auth(`Authentication attempt for: ${args.username || 'undefined'}`); try { if (!args.username || !args.password) { this.logger.auth('Authentication failed: Missing credentials'); return { content: [ { type: 'text', text: 'Authentication failed: Both username and password are required.', }, ], }; } const credentials: UserCredentials = { username: args.username, password: args.password }; const result = await this.sessionManager.authenticateUser(credentials); if (result.success) { this.logger.auth(`Authentication successful for: ${args.username}`); return { content: [ { type: 'text', text: `✅ Successfully authenticated as ${args.username}\n\nYou can now use the API tools to query your Umbrella Cost data.\n\nTry: "Show me the list of customers" or "Do I have any cost optimization recommendations?"`, }, ], }; } else { this.logger.auth(`Authentication failed for: ${args.username} - ${result.error}`); return { content: [ { type: 'text', text: `❌ Authentication failed: ${result.error}\n\nPlease verify your credentials and try again.`, }, ], }; } } catch (error: any) { this.logger.auth(`Authentication exception for: ${args.username} - ${error.message}`); return { content: [ { type: 'text', text: `❌ Authentication error: ${error.message}`, }, ], }; } } private handleSessionStatus() { this.logger.info('Session status request'); try { const stats = this.sessionManager.getStats(); const activeSessions = this.sessionManager.getActiveSessions(); if (activeSessions.length === 0) { return { content: [ { type: 'text', text: 'No active sessions. Please authenticate first using the authenticate_user tool.', }, ], }; } let output = `📊 **Active Sessions:** ${activeSessions.length}\n\n`; activeSessions.forEach(session => { output += `- ${session.username} (active)\n`; }); return { content: [ { type: 'text', text: output, }, ], }; } catch (error: any) { this.logger.error(`Session status error: ${error.message}`); return { content: [ { type: 'text', text: `Error checking session status: ${error.message}`, }, ], }; } } private handleHelp() { this.logger.info('Help request'); const helpText = `# Umbrella MCP Server Help ## Available Tools: 1. **authenticate_user** - Authenticate with your Umbrella Cost credentials 2. **session_status** - Check current authentication status 3. **api_user_management_accounts** - List MSP customer accounts 4. **api_recommendations_report** - Get cost optimization recommendations 5. **api_invoices_service_names_distinct** - Get AWS service names ## Getting Started: 1. First authenticate: authenticate_user(username="your-email", password="your-pass") 2. Check status: session_status() 3. Use API tools to get your data ## Natural Language: You can also ask questions like: - "Show me the list of customers" - "Do I have any cost optimization recommendations?" - "What AWS services am I using?"`; return { content: [ { type: 'text', text: helpText, }, ], }; } private async handleApiCall(path: string, args: any) { this.logger.info(`API call to: ${path}`); try { // Find active session const activeSessions = this.sessionManager.getActiveSessions(); if (activeSessions.length === 0) { this.logger.error('No authenticated sessions for API call'); return { content: [ { type: 'text', text: 'No authenticated users found. Please authenticate first using authenticate_user.', }, ], }; } // Use first active session (or specified user) const username = args?.username || activeSessions[0].username; const session = this.sessionManager.getUserSessionByUsername(username); if (!session) { this.logger.error(`No session found for user: ${username}`); return { content: [ { type: 'text', text: `No active session found for user: ${username}`, }, ], }; } // Remove username from API args const { username: _, ...apiArgs } = args || {}; this.logger.info(`Making API request to ${path} for user: ${username}`); const response = await session.apiClient.makeRequest(path, apiArgs); if (response.success) { this.logger.info(`API call successful: ${path}`); let output = `# API Response: ${path}\n\n`; output += `**Status:** ✅ Success\n`; output += `**User:** ${username}\n\n`; if (response.data) { if (Array.isArray(response.data)) { output += `**Results:** ${response.data.length} items\n\n`; if (path === '/user-management/accounts') { // Format customer accounts nicely output += '**Customer Accounts:**\n'; response.data.forEach((account: any, index: number) => { output += `${index + 1}. ${account.accountName || account.accountKey || 'Unknown'}\n`; output += ` - ID: ${account.accountId}\n`; output += ` - Cloud: ${this.getCloudTypeName(account.cloudTypeId)}\n\n`; }); } else { output += '```json\n'; output += JSON.stringify(response.data.slice(0, 3), null, 2); if (response.data.length > 3) { output += `\n... and ${response.data.length - 3} more items`; } output += '\n```\n'; } } else { output += '```json\n'; output += JSON.stringify(response.data, null, 2); output += '\n```\n'; } } return { content: [ { type: 'text', text: output, }, ], }; } else { this.logger.error(`API call failed: ${path} - ${response.error}`); return { content: [ { type: 'text', text: `❌ API Error: ${response.error}`, }, ], }; } } catch (error: any) { this.logger.error(`API call exception: ${path} - ${error.message}`); return { content: [ { type: 'text', text: `❌ Request Error: ${error.message}`, }, ], }; } } private getCloudTypeName(cloudTypeId: number): string { const cloudTypes: {[key: number]: string} = { 1: 'Azure', 2: 'Google Cloud', 3: 'AWS', 4: 'Multi-Cloud' }; return cloudTypes[cloudTypeId] || `Cloud ${cloudTypeId}`; } async run(): Promise<void> { this.logger.info('Starting MCP server transport'); try { const transport = new StdioServerTransport(); await this.server.connect(transport); this.logger.info('MCP Server started successfully'); console.error('🚀 Robust Umbrella MCP Server started'); console.error('📡 Ready for Claude Desktop connections'); console.error('💡 Use authenticate_user to get started'); // Graceful shutdown handlers process.on('SIGINT', () => { this.logger.info('Received SIGINT - shutting down'); this.shutdown(); process.exit(0); }); process.on('SIGTERM', () => { this.logger.info('Received SIGTERM - shutting down'); this.shutdown(); process.exit(0); }); } catch (error: any) { this.logger.error(`Failed to start server: ${error.message}`); throw error; } } shutdown(): void { this.logger.info('Shutting down server'); this.sessionManager.shutdown(); } }