Add SourceGenerator driven 'WithToolsFromAssembly' method

[dogdie233]

Is your feature request related to a problem? Please describe.
In a NativeAot project, use WithToolsFromAssembly() will cause IL2026 warning.

Describe the solution you'd like
We can collect all tool/prompt type and register them by SourceGenerator. I have made up demo

using System.Linq;
using System.Threading;

using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;

namespace ModelContextProtocol.SourceGenerator;

/// <summary>
/// An incremental source generator that creates extension methods for IMcpServerBuilder to register tools and prompts
/// from assemblies. The generator looks for classes annotated with either <see cref="T:ModelContextProtocol.Server.McpServerToolTypeAttribute" /> or 
/// <see cref="T:ModelContextProtocol.Server.McpServerPromptTypeAttribute" /> and generates corresponding registration methods.
/// </summary>
[Generator(LanguageNames.CSharp)]
public class McpToolsSourceGenerator : IIncrementalGenerator
{
    private const string ExtClassName = "SGMcpServerBuilderExtensions";
    private const string ToolsMethodName = "WithToolsFromAssemblySourceGen";
    private const string PromptsMethodName = "WithPromptsFromAssemblySourceGen";
    private const string DeclarationScript = $$"""
                                             // <auto-generated/>
                                             using Microsoft.Extensions.DependencyInjection;
                                             using System.Diagnostics.CodeAnalysis;
                                             using System.Text.Json;
                                             
                                             namespace ModelContextProtocol.Server;

                                             internal static partial class {{ExtClassName}} {
                                                 /// <summary>
                                                 /// Adds types marked with the <see cref="T:ModelContextProtocol.Server.McpServerToolTypeAttribute" /> attribute from your assembly as tools to the server. 
                                                 /// This method's behaviour is similar to <see cref="M:Microsoft.Extensions.DependencyInjection.McpServerBuilderExtensions.WithToolsFromAssembly(Microsoft.Extensions.DependencyInjection.IMcpServerBuilder,System.Reflection.Assembly,System.Text.Json.JsonSerializerOptions)" />.
                                                 /// </summary>
                                                 internal static partial IMcpServerBuilder {{ToolsMethodName}}(this IMcpServerBuilder builder, JsonSerializerOptions serializerOptions = null);
                                                 
                                                 /// <summary>
                                                 /// Adds types marked with the <see cref="T:ModelContextProtocol.Server.McpServerPromptTypeAttribute" /> attribute from your assembly as prompts to the server. 
                                                 /// This method's behaviour is similar to <see cref="M:Microsoft.Extensions.DependencyInjection.McpServerBuilderExtensions.WithPromptsFromAssembly(Microsoft.Extensions.DependencyInjection.IMcpServerBuilder,System.Reflection.Assembly,System.Text.Json.JsonSerializerOptions)" />.
                                                 /// </summary>
                                                 internal static partial IMcpServerBuilder {{PromptsMethodName}}(this IMcpServerBuilder builder, JsonSerializerOptions serializerOptions = null);
                                                 
                                                 [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access otherwise can break functionality when trimming application code", Justification = "<Pending>")]
                                                 private static void WithTool(IMcpServerBuilder builder,
                                                    [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicMethods |
                                                        DynamicallyAccessedMemberTypes.NonPublicMethods |
                                                        DynamicallyAccessedMemberTypes.PublicConstructors)] Type t,
                                                        JsonSerializerOptions serializerOptions) {
                                                     builder.WithTools([t], serializerOptions: serializerOptions);
                                                 }
                                                 
                                                 [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access otherwise can break functionality when trimming application code", Justification = "<Pending>")]
                                                 private static void WithPrompt(IMcpServerBuilder builder,
                                                    [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicMethods |
                                                        DynamicallyAccessedMemberTypes.NonPublicMethods |
                                                        DynamicallyAccessedMemberTypes.PublicConstructors)] Type t,
                                                        JsonSerializerOptions serializerOptions) {
                                                     builder.WithPrompts([t], serializerOptions: serializerOptions);
                                                 }
                                             }
                                             """;
    
    public void Initialize(IncrementalGeneratorInitializationContext context)
    {
        context.RegisterPostInitializationOutput(initializationContext =>
        {
            initializationContext.AddSource("SGMcpServerBuilderExtensions.Declaration.g.cs", DeclarationScript);
        });

        var tools = context.SyntaxProvider.ForAttributeWithMetadataName(
            "ModelContextProtocol.Server.McpServerToolTypeAttribute",
            (node, _) => node.IsKind(SyntaxKind.ClassDeclaration) || node.IsKind(SyntaxKind.RecordDeclaration),
            GetFullname).Collect();

        var prompts = context.SyntaxProvider.ForAttributeWithMetadataName(
            "ModelContextProtocol.Server.McpServerPromptTypeAttribute",
            (node, _) => node.IsKind(SyntaxKind.ClassDeclaration) || node.IsKind(SyntaxKind.RecordDeclaration),
            GetFullname).Collect();
            
        context.RegisterImplementationSourceOutput(tools, (productionContext, arr) =>
        {
            var body = string.Join("\n", 
                arr.Select(type => $"WithTool(builder, typeof({type}), serializerOptions: serializerOptions);")
            );
            productionContext.AddSource("SGMcpServerBuilderExtensions.ToolImplementation.g.cs",
                $$"""
                // <auto-generated/>
                using Microsoft.Extensions.DependencyInjection;
                using System.Diagnostics.CodeAnalysis;
                using System.Text.Json;
                
                namespace ModelContextProtocol.Server;

                internal static partial class {{ExtClassName}} {
                    internal static partial IMcpServerBuilder {{ToolsMethodName}}(this IMcpServerBuilder builder, JsonSerializerOptions serializerOptions) {
                        {{body}}
                        return builder;
                    }
                }
                """);
        });
        
        context.RegisterImplementationSourceOutput(prompts, (productionContext, arr) =>
        {
            var body = string.Join("\n",
                arr.Select(type => $"WithPrompt(builder, typeof({type}), serializerOptions: serializerOptions);")
            );
            productionContext.AddSource("SGMcpServerBuilderExtensions.PromptImplementation.g.cs",
                $$"""
                // <auto-generated/>
                using Microsoft.Extensions.DependencyInjection;
                using System.Diagnostics.CodeAnalysis;
                using System.Text.Json;
                
                namespace ModelContextProtocol.Server;

                internal static partial class {{ExtClassName}} {
                    internal static partial IMcpServerBuilder {{PromptsMethodName}}(this IMcpServerBuilder builder, JsonSerializerOptions serializerOptions) {
                        {{body}}
                        return builder;
                    }
                }
                """);
        });
        return;

        string GetFullname(GeneratorAttributeSyntaxContext syntaxContext, CancellationToken ct)
        {
            var ns = syntaxContext.TargetSymbol.ContainingNamespace;
            var name = syntaxContext.TargetSymbol.Name;
            return $"global::{(ns?.IsGlobalNamespace ?? true ? "" : ns + ".")}{name}";
        }
    }
}

To use it

builder.Services
    .AddMcpServer()
    .WithStdioServerTransport()
    // .WithToolsFromAssemblySourceGen();  // Custom json serializer options also support
    .WithToolsFromAssemblySourceGen(GeneratedJsonContext.Default.Options);

Describe alternatives you've considered
Consider adding the WithTool/WithPrompt methods to the McpServerBuilderExtensions in the SDK, and modify other methods to use them.

[eiriktsarpalis]

The simplest generation strategy I can think of is produce an interceptor that replaces WithToolsFromAssembly() with one call that chains WithTools<Tool1>().WithTools<Tool2>()...WithTools<ToolN>() where Tool i are the tools discovered by the the generator in the provided assembly.

[dogdie233]

with one call that chains WithTools<Tool1>().WithTools<Tool2>()...WithTools<ToolN>() where Tool i are the tools discovered by the the generator in the provided assembly.

That's actually a great idea, and it was my initial approach too. However, if static Tool/Prompt classes exist, it would generate code like builder.WithTools<ToolType>(), resulting in compile-time error CS0718: "Static types cannot be used as type arguments". This is why we had to implement the current slightly awkward workaround.

produce an interceptor that replaces WithToolsFromAssembly()

Regarding interceptors, if I recall correctly, they can't intercept calls to library methods? I'm not entirely certain about this, but that's my understanding.

[halter73]

Regarding interceptors, if I recall correctly, they can't intercept calls to library methods? I'm not entirely certain about this, but that's my understanding.

What do you mean by library methods? I remember there was discussion about locking interceptors down, so libraries could opt-in to making their methods interceptable, but I don't think that ever came to fruition. I also don't see any reference to such a restriction in https://github.com/dotnet/roslyn/blob/main/docs/features/interceptors.md

I could be missing something, bu taking another look at the input validation source generator @captainsafia recently wrote in dotnet/aspnetcore#60724, I don't see anything special about ValidationServiceCollectionExtensions.AddValidation in Http.Abstractions.dll that would give Http.ValidationsGenerator.dll special access to intercept it.

[dogdie233]

Okay, I stand corrected. The actual limitation is intercepting calls between precompiled methods, not the methods themselves. Really appreciate you clarifying this for me!

[stephentoub]

The simplest generation strategy I can think of is produce an interceptor that replaces WithToolsFromAssembly() with one call that chains WithTools<Tool1>().WithTools<Tool2>()...WithTools<ToolN>() where Tool i are the tools discovered by the the generator in the provided assembly.

Right now WithToolsFromAssembly uses private reflection to access tools that may not be statically nameable, e.g. a private method in another type. An intercepting source generator would either end up changing that behavior or it would not be able to do the rewrite or would be forced to use reflection in the generated code, none of which is great. If this is a direction we want to pursue, we should change the behavior of WithToolsFromAssembly to be something that can be fully replaceable by a source generator.

[eiriktsarpalis]

Right now WithToolsFromAssembly uses private reflection to access tools that may not be statically nameable, e.g. a private method in another type.

TIL -- I don't think we should support private/inaccessible tools classes using reflection just because we can. By restricting support to public/accessible members we ensure parity between the reflection and source generated worlds. This is the default in the case of STJ as well.