pfm/PostFundManagement.Application/Mcp/InAppMcpClient.cs

using PostFundManagement.Application.Abstractions.Mcp;

namespace PostFundManagement.Application.Mcp;

public class InAppMcpClient : IInAppMcpClient
{
    private static bool initialized;

    private readonly IServiceProvider serviceProvider;
    private readonly ILoggerFactory loggerFactory;

    private static readonly ConcurrentDictionary<string, ToolMethodInfo> ToolRegistry = new(StringComparer.OrdinalIgnoreCase);
    private static readonly List<McpToolDefinition> ToolDefinitions = [];
    private static readonly Lock InitLock = new();

    private static readonly JsonSerializerOptions JsonOptions = new()
    {
        PropertyNameCaseInsensitive = true
    };

    public static readonly Type[] ToolTypes =
    [
        typeof(GrantsMcpTools),
        typeof(ProjectsMcpTools),
        typeof(ExecutionMcpTools),
        typeof(GovernanceMcpTools),
        typeof(EvidenceMcpTools),
        typeof(CommunicationsMcpTools),
        typeof(UsersMcpTools)
    ];

    public InAppMcpClient(IServiceProvider serviceProvider, ILoggerFactory loggerFactory)
    {
        this.serviceProvider = serviceProvider;
        this.loggerFactory = loggerFactory;

        EnsureInitialized();
    }

    private static void EnsureInitialized()
    {
        if (initialized) return;

        lock (InitLock)
        {
            if (initialized) return;

            foreach (var toolType in ToolTypes)
            {
                var groupName = toolType.Name.Replace("McpTools", string.Empty);

                foreach (var method in toolType.GetMethods(BindingFlags.Public | BindingFlags.Instance))
                {
                    var toolAttr = method.GetCustomAttribute<McpServerToolAttribute>();

                    if (toolAttr == null) continue;

                    var toolName = !string.IsNullOrWhiteSpace(toolAttr.Name) ? toolAttr.Name : method.Name;
                    var title = !string.IsNullOrWhiteSpace(toolAttr.Title) ? toolAttr.Title : toolName;
                    var descAttr = method.GetCustomAttribute<DescriptionAttribute>();
                    var description = descAttr?.Description ?? title;

                    var parameters = new List<McpToolParameter>();
                    var methodParams = method.GetParameters();

                    foreach (var p in methodParams)
                    {
                        if (p.ParameterType == typeof(CancellationToken)) continue;

                        var pDescAttr = p.GetCustomAttribute<DescriptionAttribute>();

                        parameters.Add(new McpToolParameter(Name: p.Name ?? "param", Type: p.ParameterType.Name, Description: pDescAttr?.Description ?? string.Empty,
                            IsRequired: !p.HasDefaultValue && !IsNullable(p.ParameterType), DefaultValue: p.HasDefaultValue ? p.DefaultValue : null));
                    }

                    var toolDef = new McpToolDefinition(Name: toolName, Title: title, Description: description, Group: groupName,
                        ReadOnly: toolAttr.ReadOnly, Parameters: parameters);

                    ToolDefinitions.Add(toolDef);

                    ToolRegistry[toolName] = new ToolMethodInfo(toolType, method, methodParams);
                }
            }

            initialized = true;
        }
    }

    public IReadOnlyList<McpToolDefinition> GetAvailableTools() => ToolDefinitions;

    public async Task<McpToolExecutionResult> ExecuteToolAsync(string toolName, IReadOnlyDictionary<string, object?>? arguments = null, CancellationToken cancellationToken = default)
    {
        if (!ToolRegistry.TryGetValue(toolName, out var toolMethod))
            return McpToolExecutionResult.Fail($"MCP tool '{toolName}' is not recognized.");

        try
        {
            var targetInstance = serviceProvider.GetRequiredService(toolMethod.TargetType);
            var args = BindArguments(toolMethod.Parameters, arguments, cancellationToken);

            var invokeResult = toolMethod.Method.Invoke(targetInstance, args);

            if (invokeResult is Task task)
            {
                await task.ConfigureAwait(false);

                var resultProperty = task.GetType().GetProperty("Result");
                var value = resultProperty?.GetValue(task);

                return McpToolExecutionResult.Ok(value);
            }

            return McpToolExecutionResult.Ok(invokeResult);
        }
        catch (TargetInvocationException ex)
        {
            var inner = ex.InnerException ?? ex;

            return McpToolExecutionResult.Fail($"Execution of tool '{toolName}' failed: {inner.Message}");
        }
        catch (Exception ex)
        {
            return McpToolExecutionResult.Fail($"Execution of tool '{toolName}' failed: {ex.Message}");
        }
    }

    public async Task<McpClient> CreateProtocolClientAsync(CancellationToken cancellationToken = default)
    {
        var serverToClient = new Pipe();
        var clientToServer = new Pipe();

        var serverIn = clientToServer.Reader.AsStream();
        var serverOut = serverToClient.Writer.AsStream();

        var clientIn = serverToClient.Reader.AsStream();
        var clientOut = clientToServer.Writer.AsStream();

        var inMemoryServices = new ServiceCollection();
        inMemoryServices.AddLogging();

        foreach (var t in ToolTypes)
            inMemoryServices.AddScoped(t, _ => serviceProvider.GetRequiredService(t));

        var mcpBuilder = inMemoryServices.AddMcpServer();

        foreach (var t in ToolTypes)
            mcpBuilder.WithTools(t, (JsonSerializerOptions?)null);

        mcpBuilder.WithStreamServerTransport(serverIn, serverOut);

        var inMemoryProvider = inMemoryServices.BuildServiceProvider();
        var hostedServices = inMemoryProvider.GetServices<IHostedService>();

        foreach (var hs in hostedServices)
            _ = hs.StartAsync(cancellationToken);

        var clientTransport = new StreamClientTransport(serverInput: clientOut, serverOutput: clientIn, loggerFactory);

        return await McpClient.CreateAsync(clientTransport, cancellationToken: cancellationToken);
    }

    private static object?[] BindArguments(ParameterInfo[] parameters, IReadOnlyDictionary<string, object?>? arguments, CancellationToken cancellationToken)
    {
        var result = new object?[parameters.Length];
        var args = arguments ?? new Dictionary<string, object?>(StringComparer.OrdinalIgnoreCase);

        for (var i = 0; i < parameters.Length; i++)
        {
            var p = parameters[i];

            if (p.ParameterType == typeof(CancellationToken))
            {
                result[i] = cancellationToken;

                continue;
            }

            var matchKey = args.Keys.FirstOrDefault(k => string.Equals(k, p.Name, StringComparison.OrdinalIgnoreCase));

            if (matchKey != null && args.TryGetValue(matchKey, out var val) && val != null)
                result[i] = ConvertValue(val, p.ParameterType);
            else if (p.HasDefaultValue)
                result[i] = p.DefaultValue;
            else if (p.ParameterType == typeof(Guid) && p.Name?.Contains("userId", StringComparison.OrdinalIgnoreCase) == true)
                result[i] = SystemUser;
            else if (IsNullable(p.ParameterType))
                result[i] = null;
            else
                result[i] = p.ParameterType.IsValueType ? Activator.CreateInstance(p.ParameterType) : null;
        }

        return result;
    }

    private static object? ConvertValue(object val, Type targetType)
    {
        if (val == null) return null;

        var underlying = Nullable.GetUnderlyingType(targetType) ?? targetType;

        if (underlying.IsInstanceOfType(val)) return val;

        if (val is JsonElement jsonElement)
            return JsonSerializer.Deserialize(jsonElement.GetRawText(), targetType, JsonOptions);

        var str = val.ToString();

        if (str == null) return null;

        if (underlying == typeof(Guid)) return Guid.Parse(str);

        if (underlying.IsEnum) return Enum.Parse(underlying, str, true);

        if (underlying.IsPrimitive || underlying == typeof(decimal) || underlying == typeof(string) || underlying == typeof(DateTime))
            return Convert.ChangeType(val, underlying, CultureInfo.InvariantCulture);

        return JsonSerializer.Deserialize(JsonSerializer.Serialize(val), targetType, JsonOptions);
    }

    private static bool IsNullable(Type type) =>
        !type.IsValueType || Nullable.GetUnderlyingType(type) != null;
}