ultimate-mcp-server/dist/core/tool-registry.js

The definitive all-in-one Model Context Protocol server for AI-assisted coding across 30+ platforms

import { Logger } from "../utils/logger.js";
import { RateLimiter } from "../utils/rate-limiter.js";
export class ToolRegistry {
    tools = new Map();
    rateLimiters = new Map();
    logger;
    constructor() {
        this.logger = new Logger("ToolRegistry");
    }
    register(tool) {
        if (this.tools.has(tool.name)) {
            throw new Error(`Tool ${tool.name} is already registered`);
        }
        this.tools.set(tool.name, tool);
        // Setup rate limiter if specified
        if (tool.rateLimit) {
            this.rateLimiters.set(tool.name, new RateLimiter(tool.rateLimit.requests, tool.rateLimit.window));
        }
        this.logger.info(`Registered tool: ${tool.name}`);
    }
    unregister(name) {
        this.tools.delete(name);
        this.rateLimiters.delete(name);
        this.logger.info(`Unregistered tool: ${name}`);
    }
    async listTools() {
        return Array.from(this.tools.values()).map((tool) => {
            const result = {
                name: tool.name,
                description: tool.description,
                inputSchema: {
                    type: "object",
                    properties: {},
                },
            };
            // Handle Zod schemas
            if (tool.inputSchema) {
                // Check if it's a Zod schema (has shape property)
                if ('shape' in tool.inputSchema && tool.inputSchema.shape) {
                    const shape = tool.inputSchema.shape;
                    const properties = {};
                    const required = [];
                    for (const [key, value] of Object.entries(shape)) {
                        // Convert Zod types to JSON Schema
                        const zodType = value;
                        properties[key] = {
                            type: this.getJsonSchemaType(zodType),
                            description: zodType._def?.description
                        };
                        // Check if field is required (not optional)
                        if (!zodType.isOptional || !zodType.isOptional()) {
                            required.push(key);
                        }
                    }
                    result.inputSchema.properties = properties;
                    if (required.length > 0) {
                        result.inputSchema.required = required;
                    }
                }
                else {
                    // Fallback to direct properties assignment
                    result.inputSchema.properties = (tool.inputSchema.properties || {});
                    // Only add required if it exists and is an array
                    if (tool.inputSchema.required && Array.isArray(tool.inputSchema.required)) {
                        result.inputSchema.required = tool.inputSchema.required;
                    }
                }
            }
            return result;
        });
    }
    getJsonSchemaType(zodType) {
        if (!zodType._def)
            return 'string';
        const typeName = zodType._def.typeName;
        switch (typeName) {
            case 'ZodString':
                return 'string';
            case 'ZodNumber':
                return 'number';
            case 'ZodBoolean':
                return 'boolean';
            case 'ZodArray':
                return 'array';
            case 'ZodObject':
                return 'object';
            case 'ZodOptional':
                return this.getJsonSchemaType(zodType._def.innerType);
            default:
                return 'string';
        }
    }
    async executeTool(name, args) {
        const tool = this.tools.get(name);
        if (!tool) {
            throw new Error(`Tool ${name} not found`);
        }
        // Check rate limit
        const rateLimiter = this.rateLimiters.get(name);
        if (rateLimiter && !rateLimiter.allow()) {
            throw new Error(`Rate limit exceeded for tool ${name}`);
        }
        try {
            // Validate arguments if schema is provided
            if (tool.inputSchema) {
                this.validateArgs(args, tool.inputSchema);
            }
            // Execute tool handler
            const result = await tool.handler(args);
            this.logger.debug(`Tool ${name} executed successfully`);
            return result;
        }
        catch (error) {
            this.logger.error(`Tool ${name} execution failed:`, error);
            throw error;
        }
    }
    validateArgs(args, schema) {
        // Handle Zod schemas
        if ('shape' in schema && schema.shape) {
            // For Zod schemas, use the parse method if available
            if ('parse' in schema && typeof schema.parse === 'function') {
                try {
                    schema.parse(args);
                    return;
                }
                catch (error) {
                    throw new Error(`Validation failed: ${error.message}`);
                }
            }
            // Fallback to manual validation for Zod schemas
            const shape = schema.shape;
            for (const [key, value] of Object.entries(shape)) {
                const zodType = value;
                if (!zodType.isOptional || !zodType.isOptional()) {
                    if (!(key in args)) {
                        throw new Error(`Missing required argument: ${key}`);
                    }
                }
            }
        }
        else {
            // Basic validation for non-Zod schemas
            const required = (schema.required || []);
            if (!Array.isArray(required)) {
                return; // Skip validation if required is not an array
            }
            for (const key of required) {
                if (!(key in args)) {
                    throw new Error(`Missing required argument: ${key}`);
                }
            }
        }
    }
    getToolCount() {
        return this.tools.size;
    }
    getTool(name) {
        return this.tools.get(name);
    }
    getToolsByTag(tag) {
        return Array.from(this.tools.values()).filter((tool) => tool.tags?.includes(tag));
    }
}
//# sourceMappingURL=tool-registry.js.map