Momento/mcp-server/index.js

#!/usr/bin/env node
/**
 * Memento MCP Server - Stdio Transport (Enhanced)
 *
 * Features:
 * - Structured error handling
 * - Configuration validation
 * - Observability metrics
 * - Input validation
 * - Better logging
 *
 * Environment:
 *   DATABASE_URL   Prisma database URL
 *   USER_ID        Optional user ID filter
 *   APP_BASE_URL   Next.js app URL (default: http://localhost:3000)
 *   MCP_LOG_LEVEL  debug, info, warn, error (default: info)
 */

import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import { PrismaClient } from '@prisma/client';
import { registerTools } from './tools.js';
import config, { validateConfig, printConfig } from './config.js';
import {
  mcpError,
  mcpErrorContent,
  McpErrors,
  getErrorCategory,
  withErrorHandling,
  logError,
} from './errors.js';
import { recordRequest, recordError, recordDbQuery, getMetricsSummary } from './metrics.js';
import { validateAndSanitize } from './validation.js';

// ═══════════════════════════════════════════════════════════════
// Configuration Validation
// ═══════════════════════════════════════════════════════════════

const configErrors = validateConfig();
if (configErrors.some((e) => e.critical)) {
  console.error('❌ CRITICAL CONFIGURATION ERRORS:');
  configErrors.forEach((e) => console.error(`   ${e.key}: ${e.message}`));
  process.exit(1);
}

if (configErrors.length > 0) {
  console.warn('⚠️  Configuration warnings:');
  configErrors.forEach((e) => console.warn(`   ${e.key}: ${e.message}`));
}

// ═══════════════════════════════════════════════════════════════
// Logging
// ═══════════════════════════════════════════════════════════════

const logLevels = { debug: 0, info: 1, warn: 2, error: 3, silent: 4 };
const currentLogLevel = logLevels[config.logLevel] ?? 1;

function log(level, ...args) {
  if (logLevels[level] >= currentLogLevel) {
    const timestamp = new Date().toISOString();
    console.error(`[${timestamp}] [${level.toUpperCase()}]`, ...args);
  }
}

// ═══════════════════════════════════════════════════════════════
// Database Setup
// ═══════════════════════════════════════════════════════════════

const databaseUrl = config.databaseUrl;
if (!databaseUrl) {
  console.error('ERROR: DATABASE_URL is required');
  process.exit(1);
}

const isPostgres = databaseUrl.startsWith('postgresql://') || databaseUrl.startsWith('postgres://');

const prisma = new PrismaClient({
  datasources: { db: { url: databaseUrl } },
  ...(isPostgres
    ? {
        datasources: {
          db: {
            url: `${databaseUrl}${databaseUrl.includes('?') ? '&' : '?'}connection_limit=${config.connectionLimit}&pool_timeout=${config.poolTimeout}`,
          },
        },
      }
    : {}),
  log: config.logLevel === 'debug' ? ['query', 'info', 'warn', 'error'] : ['warn', 'error'],
});

// Wrap Prisma for metrics
const originalQuery = prisma.$queryRaw.bind(prisma);
prisma.$queryRaw = async (...args) => {
  const start = Date.now();
  try {
    const result = await originalQuery(...args);
    recordDbQuery(true, Date.now() - start);
    return result;
  } catch (error) {
    recordDbQuery(false, Date.now() - start);
    throw error;
  }
};

const appBaseUrl = config.appBaseUrl;

// ═══════════════════════════════════════════════════════════════
// MCP Server Setup with Input Validation
// ═══════════════════════════════════════════════════════════════

const server = new Server(
  { name: 'memento-mcp-server', version: '3.2.0' },
  { capabilities: { tools: {} } },
);

// Wrap tool calls with validation
const originalCallTool = server.callTool.bind(server);
server.callTool = async (request) => {
  const { name, arguments: args } = request.params;

  // Validate input
  const validation = validateAndSanitize(name, args || {});
  if (!validation.success) {
    log('warn', `Validation failed for ${name}:`, validation.errors);
    return mcpErrorContent(McpErrors.INVALID_PARAMS.code, {
      detail: 'Input validation failed',
      field: validation.errors[0]?.field,
      context: { toolName: name, errors: validation.errors },
    });
  }

  // Record tool execution
  const start = Date.now();
  try {
    const result = await originalCallTool(request);
    recordRequest(name, 'success', 'stdio', Date.now() - start);
    return result;
  } catch (error) {
    recordRequest(name, 'error', 'stdio', Date.now() - start);
    recordError(getErrorCategory(error.code || McpErrors.INTERNAL_ERROR.code), error.code || 'unknown', {
      tool: name,
    });
    throw error;
  }
};

registerTools(server, prisma, {
  userId: config.userId || null,
  appBaseUrl,
});

// ═══════════════════════════════════════════════════════════════
// Main
// ═══════════════════════════════════════════════════════════════

async function main() {
  try {
    await prisma.$queryRaw`SELECT 1`;
  } catch (error) {
    console.error('FATAL: Database connection failed:', error.message);
    process.exit(1);
  }

  const transport = new StdioServerTransport();
  await server.connect(transport);

  // Print configuration to stderr (won't interfere with stdio protocol)
  if (config.logLevel !== 'silent') {
    console.error(`
╔═══════════════════════════════════════════════════════╗
║  Memento MCP Server v3.2.0 (Enhanced)                  ║
║  Stdio Transport                                       ║
╚═══════════════════════════════════════════════════════╝

  Mode:      stdio
  Database:  ${isPostgres ? 'PostgreSQL' : 'SQLite'}
  User:      ${config.userId || 'all'}
  Log Level: ${config.logLevel}
  Tools:     29
`);
  }

  log('info', 'MCP Server ready');
}

main().catch((error) => {
  console.error('Server error:', error);
  process.exit(1);
});

// ═══════════════════════════════════════════════════════════════
// Shutdown Handler
// ═══════════════════════════════════════════════════════════════

async function shutdown() {
  log('info', 'Shutting down...');
  await prisma.$disconnect();
  process.exit(0);
}

process.on('SIGINT', shutdown);
process.on('SIGTERM', shutdown);
process.on('uncaughtException', (err) => {
  logError(log, err);
  process.exit(1);
});
process.on('unhandledRejection', (reason) => {
  log('error', 'Unhandled rejection:', reason);
  recordError(getErrorCategory(McpErrors.INTERNAL_ERROR.code), 'unhandled_rejection');
});