@boundless-oss/atlas/src/core/tool-framework.ts

Atlas - MCP Server for comprehensive startup project management

import { JSONSchema7 } from 'json-schema';
import { AtlasTool, RequestContext, ToolResult, ToolError, ToolRegistration, ToolExecutionMetadata } from './types.js';
import { SchemaValidator, getValidator, createValidationError } from './validation.js';
import { getSQLiteManager } from '../storage/sqlite-manager.js';
import { randomUUID } from 'crypto';

/**
 * Tool Registry - manages all registered tools
 */
export class ToolRegistry {
  private tools = new Map<string, AtlasTool>();
  private modules = new Map<string, ToolRegistration>();
  private validator: SchemaValidator;

  constructor() {
    this.validator = getValidator();
  }

  /**
   * Register a module with its tools
   */
  registerModule(registration: ToolRegistration): void {
    // Validate all tool schemas first
    for (const tool of registration.tools) {
      this.validateToolDefinition(tool);
    }

    // Register the module
    this.modules.set(registration.module, registration);

    // Register individual tools with module prefix
    for (const tool of registration.tools) {
      const fullName = `${registration.module}--${tool.name}`;
      this.tools.set(fullName, {
        ...tool,
        name: fullName
      });
    }

    console.error(`🔧 Registered module '${registration.module}' with ${registration.tools.length} tools`);
  }

  /**
   * Get a tool by name
   */
  getTool(name: string): AtlasTool | undefined {
    return this.tools.get(name);
  }

  /**
   * Get all registered tools
   */
  getAllTools(): AtlasTool[] {
    return Array.from(this.tools.values());
  }

  /**
   * Get tools by module
   */
  getToolsByModule(moduleName: string): AtlasTool[] {
    const registration = this.modules.get(moduleName);
    if (!registration) return [];

    return registration.tools.map(tool => ({
      ...tool,
      name: `${moduleName}.${tool.name}`
    }));
  }

  /**
   * Get all modules
   */
  getModules(): string[] {
    return Array.from(this.modules.keys());
  }

  /**
   * Get tool manifest for MCP
   */
  getToolManifest(): Array<{
    name: string;
    description: string;
    inputSchema: JSONSchema7;
    category?: string;
    readOnly?: boolean;
  }> {
    return this.getAllTools().map(tool => ({
      name: tool.name,
      description: tool.description,
      inputSchema: tool.inputSchema,
      category: tool.category,
      readOnly: tool.readOnly
    }));
  }

  /**
   * Get tool count
   */
  getToolCount(): number {
    return this.tools.size;
  }

  /**
   * Get module count
   */
  getModuleCount(): number {
    return this.modules.size;
  }

  /**
   * Validate a tool definition
   */
  private validateToolDefinition(tool: AtlasTool): void {
    if (!tool.name || typeof tool.name !== 'string') {
      throw new Error('Tool name is required and must be a string');
    }

    if (!tool.description || typeof tool.description !== 'string') {
      throw new Error('Tool description is required and must be a string');
    }

    if (!tool.inputSchema || typeof tool.inputSchema !== 'object') {
      throw new Error('Tool inputSchema is required and must be a valid JSON schema');
    }

    if (typeof tool.execute !== 'function') {
      throw new Error('Tool execute function is required');
    }

    // Validate the schema itself
    try {
      this.validator.createValidator(tool.inputSchema);
    } catch (error) {
      throw new Error(`Invalid input schema for tool '${tool.name}': ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }
}

/**
 * Tool executor - handles tool execution with validation and error handling
 */
export class ToolExecutor {
  private registry: ToolRegistry;
  private validator: SchemaValidator;

  constructor(registry: ToolRegistry) {
    this.registry = registry;
    this.validator = getValidator();
  }

  /**
   * Execute a tool with validation and error handling
   */
  async execute(
    toolName: string,
    input: any,
    context: RequestContext
  ): Promise<ToolResult> {
    const startTime = Date.now();

    try {
      // Get the tool
      const tool = this.registry.getTool(toolName);
      if (!tool) {
        return this.createErrorResult({
          code: 'TOOL_NOT_FOUND',
          message: `Tool '${toolName}' not found`,
          details: { toolName, availableTools: this.registry.getAllTools().map(t => t.name) },
          suggestions: [
            'Check the tool name for typos',
            'Use list_tools to see available tools',
            'Verify the tool module is loaded'
          ],
          recoverable: false,
          category: 'resource'
        }, context.requestId);
      }

      // Validate input
      const validationResult = this.validator.validate(tool.inputSchema, input);
      if (!validationResult.valid) {
        return this.createErrorResult(
          createValidationError(validationResult.errors),
          context.requestId
        );
      }

      // Record execution start
      await this.recordExecution({
        toolName,
        requestId: context.requestId,
        userId: context.userId,
        startTime,
        success: false // Will be updated after execution
      });

      // Execute the tool
      const result = await tool.execute(input, context);

      // Calculate execution time
      const executionTime = Date.now() - startTime;

      // Update execution record
      await this.recordExecution({
        toolName,
        requestId: context.requestId,
        userId: context.userId,
        startTime,
        endTime: Date.now(),
        executionTime,
        success: result.success,
        errorCode: result.error?.code
      });

      // Add metadata
      if (!result.metadata) {
        result.metadata = {};
      }
      result.metadata.executionTime = executionTime;
      result.metadata.requestId = context.requestId;

      return result;

    } catch (error) {
      const executionTime = Date.now() - startTime;

      // Record failed execution
      await this.recordExecution({
        toolName,
        requestId: context.requestId,
        userId: context.userId,
        startTime,
        endTime: Date.now(),
        executionTime,
        success: false,
        errorCode: 'EXECUTION_ERROR'
      });

      return this.createErrorResult({
        code: 'EXECUTION_ERROR',
        message: `Tool execution failed: ${error instanceof Error ? error.message : 'Unknown error'}`,
        details: {
          toolName,
          error: error instanceof Error ? error.stack : String(error)
        },
        suggestions: [
          'Check tool implementation for bugs',
          'Verify input parameters are correct',
          'Check system resources and dependencies',
          'Review logs for more details'
        ],
        recoverable: true,
        category: 'execution'
      }, context.requestId, executionTime);
    }
  }

  /**
   * Create a standardized error result
   */
  private createErrorResult(
    error: ToolError,
    requestId: string,
    executionTime?: number
  ): ToolResult {
    return {
      success: false,
      error,
      metadata: {
        requestId,
        executionTime
      }
    };
  }

  /**
   * Record tool execution metrics
   */
  private async recordExecution(metadata: ToolExecutionMetadata): Promise<void> {
    try {
      const db = getSQLiteManager();

      if (metadata.endTime) {
        // Update existing record
        await db.run(
          `UPDATE performance_metrics 
           SET execution_time = ?, success = ?, error_details = ? 
           WHERE request_id = ?`,
          [
            metadata.executionTime || 0,
            metadata.success ? 1 : 0,
            metadata.errorCode || null,
            metadata.requestId
          ]
        );
      } else {
        // Insert new record
        await db.run(
          `INSERT INTO performance_metrics 
           (id, tool_name, execution_time, success, user_id, request_id, created_at) 
           VALUES (?, ?, ?, ?, ?, ?, ?)`,
          [
            randomUUID(),
            metadata.toolName,
            0, // Will be updated later
            0, // Will be updated later
            metadata.userId || null,
            metadata.requestId,
            metadata.startTime
          ]
        );
      }
    } catch (error) {
      console.error('Failed to record execution metrics:', error);
      // Don't fail the tool execution for metrics recording errors
    }
  }
}

/**
 * Global registry instance
 */
let globalRegistry: ToolRegistry | null = null;
let globalExecutor: ToolExecutor | null = null;

export function getToolRegistry(): ToolRegistry {
  if (!globalRegistry) {
    globalRegistry = new ToolRegistry();
  }
  return globalRegistry;
}

export function getToolExecutor(): ToolExecutor {
  if (!globalExecutor) {
    globalExecutor = new ToolExecutor(getToolRegistry());
  }
  return globalExecutor;
}

/**
 * Utility function to create a tool definition with common patterns
 */
export function createTool<TInput, TOutput>(config: {
  name: string;
  description: string;
  inputSchema: JSONSchema7;
  outputSchema?: JSONSchema7;
  category?: string;
  requiresApproval?: boolean;
  readOnly?: boolean;
  execute: (input: TInput, context: RequestContext) => Promise<ToolResult<TOutput>>;
}): AtlasTool<TInput, TOutput> {
  return {
    name: config.name,
    description: config.description,
    inputSchema: config.inputSchema,
    outputSchema: config.outputSchema,
    category: config.category,
    requiresApproval: config.requiresApproval || false,
    readOnly: config.readOnly || false,
    execute: config.execute
  };
}

/**
 * Utility to create success results
 */
export function createSuccessResult<T>(data: T, metadata?: any): ToolResult<T> {
  return {
    success: true,
    data,
    metadata
  };
}

/**
 * Utility to create error results
 */
export function createErrorResult(error: Partial<ToolError>): ToolResult {
  return {
    success: false,
    error: {
      code: error.code || 'UNKNOWN_ERROR',
      message: error.message || 'An unknown error occurred',
      details: error.details || {},
      suggestions: error.suggestions || ['Please try again'],
      recoverable: error.recoverable !== false,
      category: error.category || 'execution'
    }
  };
}