langterm/lib/mcp.js

Secure CLI tool that translates natural language to shell commands using local AI models via Ollama, with project memory system, reusable command templates (hooks), MCP (Model Context Protocol) support, and dangerous command detection

import { Client } from '@modelcontextprotocol/sdk/client/index.js';
import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js';
import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js';
import chalk from 'chalk';
import { analyzeIntent, INTENT_TYPES, getIntentDescription } from './intent.js';
import { prompt } from './security.js';

/**
 * MCP client manager for connecting to Model Context Protocol servers
 */
export class MCPManager {
  constructor() {
    this.clients = new Map();
    this.isInitialized = false;
  }

  /**
   * Initialize MCP connections based on configuration
   * @param {object} mcpConfig - MCP server configurations
   */
  async initialize(mcpConfig) {
    if (!mcpConfig || !mcpConfig.servers) {
      return;
    }

    console.log(chalk.gray('🔗 Initializing MCP connections...'));

    for (const [serverName, config] of Object.entries(mcpConfig.servers)) {
      try {
        await this.connectToServer(serverName, config);
      } catch (error) {
        console.log(chalk.yellow(`⚠️  Failed to connect to MCP server '${serverName}': ${error.message}`));
      }
    }

    this.isInitialized = true;
  }

  /**
   * Connect to a single MCP server
   * @param {string} name - Server name
   * @param {object} config - Server configuration
   */
  async connectToServer(name, config) {
    let transport;
    let client;

    try {
      if (config.type === 'stdio') {
        // For stdio transport (local commands)
        // Debug logging
        if (process.env.DEBUG) {
          console.log(chalk.gray(`Setting up MCP server: ${config.command} ${(config.args || []).join(' ')}`));
        }
        
        transport = new StdioClientTransport({
          command: config.command,
          args: config.args || [],
          env: { ...process.env, ...(config.env || {}) }
        });
      } else if (config.type === 'sse') {
        // For Server-Sent Events transport (HTTP)
        transport = new SSEClientTransport(new URL(config.url));
      } else {
        throw new Error(`Unsupported transport type: ${config.type}`);
      }

      client = new Client({
        name: 'langterm',
        version: '1.0.0'
      }, {
        capabilities: {
          sampling: {}
        }
      });

      await client.connect(transport);
      
      this.clients.set(name, {
        client,
        transport,
        config
      });

      console.log(chalk.green(`✅ Connected to MCP server '${name}'`));
    } catch (error) {
      if (client) {
        try {
          await client.close();
        } catch (closeError) {
          // Ignore close errors
        }
      }
      throw error;
    }
  }

  /**
   * Get available tools from all connected MCP servers
   * @returns {Promise<Array>} Array of available tools
   */
  async getAvailableTools() {
    const allTools = [];

    for (const [serverName, { client }] of this.clients) {
      try {
        const result = await client.listTools();
        const tools = result.tools || [];
        
        for (const tool of tools) {
          allTools.push({
            ...tool,
            serverName,
            fullName: `${serverName}:${tool.name}`
          });
        }
      } catch (error) {
        console.log(chalk.yellow(`⚠️  Failed to list tools from '${serverName}': ${error.message}`));
      }
    }

    return allTools;
  }

  /**
   * Get available resources from all connected MCP servers
   * @returns {Promise<Array>} Array of available resources
   */
  async getAvailableResources() {
    const allResources = [];

    for (const [serverName, { client }] of this.clients) {
      try {
        const result = await client.listResources();
        const resources = result.resources || [];
        
        for (const resource of resources) {
          allResources.push({
            ...resource,
            serverName,
            fullName: `${serverName}:${resource.uri}`
          });
        }
      } catch (error) {
        console.log(chalk.yellow(`⚠️  Failed to list resources from '${serverName}': ${error.message}`));
      }
    }

    return allResources;
  }

  /**
   * Get available prompts from all connected MCP servers
   * @returns {Promise<Array>} Array of available prompts
   */
  async getAvailablePrompts() {
    const allPrompts = [];

    for (const [serverName, { client }] of this.clients) {
      try {
        const result = await client.listPrompts();
        const prompts = result.prompts || [];
        
        for (const prompt of prompts) {
          allPrompts.push({
            ...prompt,
            serverName,
            fullName: `${serverName}:${prompt.name}`
          });
        }
      } catch (error) {
        console.log(chalk.yellow(`⚠️  Failed to list prompts from '${serverName}': ${error.message}`));
      }
    }

    return allPrompts;
  }

  /**
   * Call a tool on a specific MCP server
   * @param {string} serverName - Name of the MCP server
   * @param {string} toolName - Name of the tool to call
   * @param {object} args - Arguments for the tool
   * @returns {Promise<object>} Tool execution result
   */
  async callTool(serverName, toolName, args = {}) {
    const serverConnection = this.clients.get(serverName);
    if (!serverConnection) {
      throw new Error(`MCP server '${serverName}' not connected`);
    }

    try {
      const result = await serverConnection.client.callTool({
        name: toolName,
        arguments: args
      });
      return result;
    } catch (error) {
      throw new Error(`Failed to call tool '${toolName}' on server '${serverName}': ${error.message}`);
    }
  }

  /**
   * Read a resource from a specific MCP server
   * @param {string} serverName - Name of the MCP server
   * @param {string} uri - URI of the resource to read
   * @returns {Promise<object>} Resource content
   */
  async readResource(serverName, uri) {
    const serverConnection = this.clients.get(serverName);
    if (!serverConnection) {
      throw new Error(`MCP server '${serverName}' not connected`);
    }

    try {
      const result = await serverConnection.client.readResource({ uri });
      return result;
    } catch (error) {
      throw new Error(`Failed to read resource '${uri}' from server '${serverName}': ${error.message}`);
    }
  }

  /**
   * Get enhanced context for command generation using MCP resources
   * @param {string} userInput - User's natural language input
   * @returns {Promise<string>} Enhanced context string
   */
  async getEnhancedContext(userInput) {
    if (!this.isInitialized || this.clients.size === 0) {
      return '';
    }

    try {
      const contextParts = [];
      
      // Get relevant resources that might provide context
      const resources = await this.getAvailableResources();
      const relevantResources = resources.filter(resource => 
        this.isResourceRelevant(resource, userInput)
      );

      // Read relevant resources (limit to prevent overwhelming the context)
      const maxResources = 3;
      for (let i = 0; i < Math.min(relevantResources.length, maxResources); i++) {
        const resource = relevantResources[i];
        try {
          const content = await this.readResource(resource.serverName, resource.uri);
          if (content && content.contents) {
            contextParts.push(`Context from ${resource.fullName}: ${JSON.stringify(content.contents)}`);
          }
        } catch (error) {
          // Skip failed resources
          continue;
        }
      }

      // Get available tools that might be relevant
      const tools = await this.getAvailableTools();
      if (tools.length > 0) {
        const toolList = tools.map(tool => `${tool.fullName}: ${tool.description || 'No description'}`).join(', ');
        contextParts.push(`Available MCP tools: ${toolList}`);
      }

      return contextParts.length > 0 ? '\n\nMCP Context:\n' + contextParts.join('\n') : '';
    } catch (error) {
      console.log(chalk.yellow(`⚠️  Failed to get MCP context: ${error.message}`));
      return '';
    }
  }

  /**
   * Check if a resource is relevant to the user input
   * @param {object} resource - MCP resource
   * @param {string} userInput - User's input
   * @returns {boolean} Whether the resource is relevant
   */
  isResourceRelevant(resource, userInput) {
    const input = userInput.toLowerCase();
    const resourceName = (resource.name || resource.uri).toLowerCase();
    const resourceDesc = (resource.description || '').toLowerCase();
    
    // Simple relevance check - can be enhanced with more sophisticated matching
    return input.includes('file') && resourceName.includes('file') ||
           input.includes('git') && resourceName.includes('git') ||
           input.includes('project') && resourceName.includes('project') ||
           resourceDesc.includes('current') ||
           resourceDesc.includes('workspace');
  }

  /**
   * Get connection status for all MCP servers
   * @returns {Array} Array of server connection statuses
   */
  getConnectionStatus() {
    const status = [];
    for (const [serverName, connection] of this.clients) {
      status.push({
        name: serverName,
        connected: true,
        type: connection.config.type,
        config: connection.config
      });
    }
    return status;
  }

  /**
   * Analyze user intent and determine the best approach
   * @param {string} userInput - User's natural language input
   * @returns {Promise<object>} Intent analysis with recommendations
   */
  async analyzeUserIntent(userInput) {
    if (!this.isInitialized) {
      return {
        intent: INTENT_TYPES.TERMINAL_COMMAND,
        confidence: 1.0,
        reasoning: ["MCP not initialized, defaulting to terminal command"],
        suggestedTools: [],
        availableTools: []
      };
    }

    try {
      const availableTools = await this.getAvailableTools();
      const analysis = analyzeIntent(userInput, availableTools);
      
      return {
        ...analysis,
        availableTools
      };
    } catch (error) {
      console.log(chalk.yellow(`⚠️  Failed to analyze intent: ${error.message}`));
      return {
        intent: INTENT_TYPES.TERMINAL_COMMAND,
        confidence: 0.5,
        reasoning: ["Intent analysis failed, defaulting to terminal command"],
        suggestedTools: [],
        availableTools: []
      };
    }
  }

  /**
   * Execute the determined intent (MCP tool, terminal command, or hybrid)
   * @param {string} userInput - User's input
   * @param {object} intentAnalysis - Result from analyzeUserIntent
   * @returns {Promise<object>} Execution result
   */
  async executeIntent(userInput, intentAnalysis) {
    const { intent, suggestedTools, confidence } = intentAnalysis;

    // Only show intent analysis if we actually have MCP tools
    if (intentAnalysis.availableTools && intentAnalysis.availableTools.length > 0) {
      console.log(chalk.gray(`🎯 Intent: ${getIntentDescription(intent)} (confidence: ${(confidence * 100).toFixed(0)}%)`));
    } else if (intentAnalysis.availableTools && intentAnalysis.availableTools.length === 0) {
      console.log(chalk.gray('No MCP tools connected, using terminal command generation.'));
    }
    
    switch (intent) {
      case INTENT_TYPES.MCP_TOOL_DIRECT:
        return await this.executeMCPToolDirect(userInput, suggestedTools);
        
      case INTENT_TYPES.HYBRID:
        return await this.executeHybridApproach(userInput, suggestedTools);
        
      case INTENT_TYPES.AMBIGUOUS:
        return await this.handleAmbiguousIntent(userInput, intentAnalysis);
        
      case INTENT_TYPES.TERMINAL_COMMAND:
      default:
        return { 
          type: 'terminal_command', 
          mcpContext: await this.getEnhancedContext(userInput),
          message: 'Proceeding with terminal command generation'
        };
    }
  }

  /**
   * Execute MCP tool directly with user confirmation
   */
  async executeMCPToolDirect(userInput, suggestedTools) {
    if (suggestedTools.length === 0) {
      console.log(chalk.yellow('⚠️  No suitable MCP tools found for this request.'));
      return { 
        type: 'fallback_to_terminal', 
        mcpContext: await this.getEnhancedContext(userInput),
        message: 'Falling back to terminal command generation'
      };
    }

    // Show user what tools would be used
    console.log(chalk.cyan('\n🔧 Suggested MCP tools:'));
    suggestedTools.forEach((tool, index) => {
      console.log(chalk.gray(`  ${index + 1}. ${tool.fullName}: ${tool.description || 'No description'}`));
    });

    const shouldExecute = await this.confirmMCPToolExecution(userInput, suggestedTools);
    
    if (!shouldExecute) {
      console.log(chalk.yellow('🔄 Switching to terminal command generation...'));
      return { 
        type: 'fallback_to_terminal', 
        mcpContext: await this.getEnhancedContext(userInput),
        message: 'User declined MCP tool execution, using terminal command'
      };
    }

    try {
      // For simplicity, use the first suggested tool
      const tool = suggestedTools[0];
      console.log(chalk.cyan(`🔧 Executing MCP tool: ${tool.fullName}`));
      
      // Generate appropriate arguments for the tool based on user input
      const args = this.generateToolArguments(userInput, tool);
      const result = await this.callTool(tool.serverName, tool.name, args);
      
      console.log(chalk.green('✅ MCP tool execution completed:'));
      console.log(JSON.stringify(result, null, 2));
      
      return {
        type: 'mcp_tool_executed',
        result,
        tool: tool.fullName,
        message: 'MCP tool executed successfully'
      };
    } catch (error) {
      console.log(chalk.red(`❌ MCP tool execution failed: ${error.message}`));
      return { 
        type: 'fallback_to_terminal', 
        mcpContext: await this.getEnhancedContext(userInput),
        message: 'MCP tool failed, falling back to terminal command'
      };
    }
  }

  /**
   * Execute hybrid approach (MCP + terminal command)
   */
  async executeHybridApproach(userInput, suggestedTools) {
    console.log(chalk.magenta('🔀 Using hybrid approach: MCP tools + terminal command'));
    
    // First try to get enhanced context from MCP tools
    let mcpData = null;
    if (suggestedTools.length > 0) {
      try {
        const shouldUseMCP = await this.confirmMCPToolExecution(
          userInput, 
          suggestedTools, 
          'Use MCP tools to gather additional context?'
        );
        
        if (shouldUseMCP) {
          const tool = suggestedTools[0];
          const args = this.generateToolArguments(userInput, tool);
          mcpData = await this.callTool(tool.serverName, tool.name, args);
          console.log(chalk.green('✅ Context gathered from MCP tools'));
        }
      } catch (error) {
        console.log(chalk.yellow(`⚠️  MCP context gathering failed: ${error.message}`));
      }
    }
    
    // Enhanced context includes both regular MCP context and tool data
    const enhancedContext = await this.getEnhancedContext(userInput);
    const hybridContext = mcpData 
      ? `${enhancedContext}\n\nMCP Tool Data: ${JSON.stringify(mcpData)}`
      : enhancedContext;
    
    return {
      type: 'hybrid_execution',
      mcpContext: hybridContext,
      mcpData,
      message: 'Using hybrid approach with enhanced context'
    };
  }

  /**
   * Handle ambiguous intent by asking user for clarification
   */
  async handleAmbiguousIntent(userInput, intentAnalysis) {
    // If no MCP tools are available, just return terminal command
    if (!intentAnalysis.availableTools || intentAnalysis.availableTools.length === 0) {
      return { 
        type: 'terminal_command', 
        mcpContext: '',
        message: 'No MCP tools available, using terminal command'
      };
    }
    
    console.log(chalk.yellow('\n🤔 Ambiguous request detected.'));
    console.log(chalk.gray('Reasoning:'), intentAnalysis.reasoning);
    
    if (intentAnalysis.suggestedTools.length > 0) {
      console.log(chalk.cyan('\n💡 Available options:'));
      console.log(chalk.gray('1. Use MCP tools for data access'));
      console.log(chalk.gray('2. Generate terminal command'));
      console.log(chalk.gray('3. Use both (hybrid approach)'));
      
      const choice = await prompt('Choose approach (1-3, or press Enter for terminal command): ');
      
      switch (choice) {
        case '1':
          return await this.executeMCPToolDirect(userInput, intentAnalysis.suggestedTools);
        case '3':
          return await this.executeHybridApproach(userInput, intentAnalysis.suggestedTools);
        case '2':
        default:
          return { 
            type: 'terminal_command', 
            mcpContext: await this.getEnhancedContext(userInput),
            message: 'Using terminal command generation'
          };
      }
    } else {
      console.log(chalk.gray('No MCP tools available, using terminal command.'));
      return { 
        type: 'terminal_command', 
        mcpContext: await this.getEnhancedContext(userInput),
        message: 'No MCP tools available, using terminal command'
      };
    }
  }

  /**
   * Confirm MCP tool execution with user
   */
  async confirmMCPToolExecution(userInput, tools, customMessage = null) {
    const message = customMessage || 'Execute MCP tools for this request?';
    console.log(chalk.yellow(`\n❓ ${message}`));
    console.log(chalk.gray(`Request: "${userInput}"`));
    
    const response = await prompt('Proceed with MCP tools? (y/n): ');
    return response.toLowerCase() === 'y' || response.toLowerCase() === 'yes';
  }

  /**
   * Generate appropriate arguments for MCP tool based on user input
   */
  generateToolArguments(userInput, tool) {
    // Simple argument generation - can be enhanced with more sophisticated logic
    const args = {};
    
    // Debug logging
    if (process.env.DEBUG) {
      console.log(chalk.gray(`Generating arguments for tool: ${tool.name}, input: "${userInput}"`));
    }
    
    // Handle directory listing tools
    if (tool.name.includes('list_directory') || tool.name.includes('directory_tree')) {
      // Check if user specified a specific directory (but not "this directory")
      const dirMatch = userInput.match(/(?:in|from|of)\s+(?!this\s|current\s)([^\s]+(?:\/[^\s]*)?)/i);
      if (dirMatch) {
        args.path = dirMatch[1];
      } else {
        // Default to current directory for any list commands
        args.path = process.cwd();
      }
    }
    
    // Common patterns for different tool types
    if (tool.name.includes('read') || tool.name.includes('get')) {
      // First try to match common files like package.json
      const commonFileMatch = userInput.match(/(package\.json|README\.md|Dockerfile|Makefile|\.gitignore|\.env)/i);
      if (commonFileMatch) {
        args.path = commonFileMatch[1];
      } else {
        // Extract file paths with extensions
        const pathMatch = userInput.match(/([a-zA-Z0-9_\-./]+\.(js|jsx|ts|tsx|py|txt|json|md|yml|yaml|config|env|html|css|xml|log|ini|conf))/i);
        if (pathMatch) {
          args.path = pathMatch[1];
        }
      }
      
      const urlMatch = userInput.match(/(https?:\/\/[^\s]+)/i);
      if (urlMatch) {
        args.url = urlMatch[1];
      }
    }
    
    if (tool.name.includes('search') || tool.name.includes('find')) {
      // Extract search terms
      const searchMatch = userInput.match(/(?:search|find)\s+(?:for\s+)?([^\s]+)/i);
      if (searchMatch) {
        args.query = searchMatch[1];
      }
    }
    
    if (process.env.DEBUG) {
      console.log(chalk.gray(`Generated args:`, JSON.stringify(args, null, 2)));
    }
    
    return args;
  }

  /**
   * Close all MCP connections
   */
  async disconnect() {
    for (const [serverName, { client }] of this.clients) {
      try {
        await client.close();
        console.log(chalk.gray(`🔌 Disconnected from MCP server '${serverName}'`));
      } catch (error) {
        console.log(chalk.yellow(`⚠️  Error disconnecting from '${serverName}': ${error.message}`));
      }
    }
    this.clients.clear();
    this.isInitialized = false;
  }
}

/**
 * Create and return a singleton MCP manager instance
 */
export const mcpManager = new MCPManager();