Momento/mcp-server/tool-handlers.js

/**
 * Memento MCP Server - Enhanced Tool Handler Wrapper
 *
 * Wraps tool handlers with error handling, metrics recording, and validation.
 * Import this in tools.js to wrap the CallToolRequestSchema handler.
 */

import { McpErrors, getErrorCategory, mcpErrorContent, logError } from './errors.js';
import { recordRequest, recordError, recordToolExecution } from './metrics.js';

/**
 * Wrap a tool handler with error handling and metrics
 *
 * @param {string} toolName - Name of the tool
 * @param {Function} handler - The actual tool handler function
 * @param {object} options - Options
 * @returns {Function} Wrapped handler
 */
export function wrapToolHandler(toolName, handler, options = {}) {
  const { timeoutMs = 60000 } = options;

  return async (args, context) => {
    const start = Date.now();

    // Create timeout promise
    const timeout = new Promise((_, reject) => {
      setTimeout(() => {
        reject(new Error(`Tool ${toolName} timed out after ${timeoutMs}ms`));
      }, timeoutMs);
    });

    try {
      // Race between handler and timeout
      const result = await Promise.race([
        handler(args, context),
        timeout,
      ]);

      // Record success
      const duration = Date.now() - start;
      recordToolExecution(toolName, true, duration);
      recordRequest(toolName, 'success', 'mcp', duration);

      return result;
    } catch (error) {
      const duration = Date.now() - start;

      // Determine error type
      let errorCode = McpErrors.INTERNAL_ERROR.code;
      let errorMessage = error.message || 'An unexpected error occurred';

      if (error.code === 'P2025') {
        errorCode = McpErrors.NOT_FOUND.code;
        errorMessage = 'Record not found';
      } else if (error.code?.startsWith('P')) {
        errorCode = McpErrors.DATABASE_ERROR.code;
        errorMessage = 'Database operation failed';
      } else if (error.message?.includes('timeout')) {
        errorCode = McpErrors.TIMEOUT.code;
      }

      // Record error
      recordError(getErrorCategory(errorCode), errorCode, { tool: toolName });
      recordToolExecution(toolName, false, duration);
      recordRequest(toolName, 'error', 'mcp', duration);

      // Log error
      logError(console, error, { tool: toolName });

      // Return error response
      return mcpErrorContent(errorCode, {
        detail: errorMessage,
        context: { tool: toolName, duration },
        cause: error,
      });
    }
  };
}

/**
 * Create a tool handler map with wrapped handlers
 *
 * @param {object} handlers - Map of tool name to handler function
 * @param {object} options - Options for wrapping
 * @returns {object} Map of wrapped handlers
 */
export function createToolHandlerMap(handlers, options = {}) {
  const wrapped = {};
  for (const [name, handler] of Object.entries(handlers)) {
    wrapped[name] = wrapToolHandler(name, handler, options);
  }
  return wrapped;
}

/**
 * Execute a tool by name with proper error handling
 *
 * @param {string} toolName - Name of the tool to execute
 * @param {object} handlers - Map of tool handlers
 * @param {object} args - Tool arguments
 * @param {object} context - Execution context (prisma, userId, etc.)
 * @param {object} options - Options
 * @returns {Promise<object>} Tool result
 */
export async function executeTool(toolName, handlers, args, context, options = {}) {
  const handler = handlers[toolName];

  if (!handler) {
    return mcpErrorContent(McpErrors.NOT_FOUND.code, {
      detail: `Tool not found: ${toolName}`,
      context: { availableTools: Object.keys(handlers) },
    });
  }

  return wrapToolHandler(toolName, handler, options)(args, context);
}

/**
 * Batch execute multiple tools
 *
 * @param {Array} operations - Array of { tool, args } objects
 * @param {object} handlers - Map of tool handlers
 * @param {object} context - Execution context
 * @param {object} options - Options
 * @returns {Promise<Array>} Array of results
 */
export async function executeBatch(operations, handlers, context, options = {}) {
  const results = [];
  const { continueOnError = true } = options;

  for (const op of operations) {
    try {
      const result = await executeTool(op.tool, handlers, op.args || {}, context, options);
      results.push({ tool: op.tool, success: true, result });
    } catch (error) {
      results.push({
        tool: op.tool,
        success: false,
        error: error.message || 'Execution failed',
      });
      if (!continueOnError) {
        break;
      }
    }
  }

  return results;
}

export default {
  wrapToolHandler,
  createToolHandlerMap,
  executeTool,
  executeBatch,
};