ColaFlow/colaflow-api/src/Modules/Mcp/ColaFlow.Modules.Mcp.Application/Services/McpToolRegistry.cs

using ColaFlow.Modules.Mcp.Contracts.Tools;
using ColaFlow.Modules.Mcp.Domain.Exceptions;
using Microsoft.Extensions.Logging;

namespace ColaFlow.Modules.Mcp.Application.Services;

/// <summary>
/// Registry for MCP Tools with auto-discovery
/// Uses constructor injection to register all IMcpTool implementations
/// </summary>
public class McpToolRegistry : IMcpToolRegistry
{
    private readonly Dictionary<string, IMcpTool> _tools;
    private readonly ILogger<McpToolRegistry> _logger;

    public McpToolRegistry(
        IEnumerable<IMcpTool> tools,
        ILogger<McpToolRegistry> _logger)
    {
        this._logger = _logger ?? throw new ArgumentNullException(nameof(_logger));

        // Auto-discover and register all tools
        _tools = new Dictionary<string, IMcpTool>(StringComparer.OrdinalIgnoreCase);

        foreach (var tool in tools)
        {
            if (_tools.ContainsKey(tool.Name))
            {
                this._logger.LogWarning(
                    "Duplicate tool name detected: {ToolName}. Skipping duplicate registration.",
                    tool.Name);
                continue;
            }

            _tools[tool.Name] = tool;
            this._logger.LogInformation(
                "Registered MCP Tool: {ToolName} - {Description}",
                tool.Name, tool.Description);
        }

        this._logger.LogInformation(
            "McpToolRegistry initialized with {Count} tools",
            _tools.Count);
    }

    public IEnumerable<IMcpTool> GetAllTools()
    {
        return _tools.Values;
    }

    public IMcpTool? GetTool(string toolName)
    {
        if (string.IsNullOrWhiteSpace(toolName))
            return null;

        _tools.TryGetValue(toolName, out var tool);
        return tool;
    }

    public bool HasTool(string toolName)
    {
        if (string.IsNullOrWhiteSpace(toolName))
            return false;

        return _tools.ContainsKey(toolName);
    }

    public async Task<McpToolResult> ExecuteToolAsync(
        McpToolCall toolCall,
        CancellationToken cancellationToken = default)
    {
        if (toolCall == null)
            throw new ArgumentNullException(nameof(toolCall));

        if (string.IsNullOrWhiteSpace(toolCall.Name))
            throw new McpInvalidParamsException("Tool name cannot be empty");

        // Get tool
        var tool = GetTool(toolCall.Name);
        if (tool == null)
        {
            _logger.LogWarning(
                "Tool not found: {ToolName}. Available tools: {AvailableTools}",
                toolCall.Name, string.Join(", ", _tools.Keys));

            throw new McpNotFoundException("Tool", toolCall.Name);
        }

        _logger.LogInformation(
            "Executing MCP Tool: {ToolName}",
            toolCall.Name);

        try
        {
            // Execute tool
            var result = await tool.ExecuteAsync(toolCall, cancellationToken);

            _logger.LogInformation(
                "MCP Tool executed successfully: {ToolName}, IsError={IsError}",
                toolCall.Name, result.IsError);

            return result;
        }
        catch (Exception ex)
        {
            _logger.LogError(ex,
                "Error executing MCP Tool: {ToolName}",
                toolCall.Name);

            // Return error result
            return new McpToolResult
            {
                Content = new[]
                {
                    new McpToolContent
                    {
                        Type = "text",
                        Text = $"Error executing tool '{toolCall.Name}': {ex.Message}"
                    }
                },
                IsError = true
            };
        }
    }
}