agent-framework/dotnet/src/Microsoft.Agents.AI.Workflows.Generators/Generation/SourceBuilder.cs

// Copyright (c) Microsoft. All rights reserved.

using System.Collections.Generic;
using System.Linq;
using System.Text;
using Microsoft.Agents.AI.Workflows.Generators.Models;

namespace Microsoft.Agents.AI.Workflows.Generators.Generation;

/// <summary>
/// Generates source code for executor route configuration.
/// </summary>
/// <remarks>
/// This builder produces a partial class file that overrides <c>ConfigureRoutes</c> to register
/// handlers discovered via [MessageHandler] attributes. It may also generate <c>ConfigureSentTypes</c>
/// and <c>ConfigureYieldTypes</c> overrides when [SendsMessage] or [YieldsOutput] attributes are present.
/// </remarks>
internal static class SourceBuilder
{
    internal const string IndentUnit = "    ";

    /// <summary>
    /// Generates the complete source file for an executor's generated partial class.
    /// </summary>
    /// <param name="info">The analyzed executor information containing class metadata and handler details.</param>
    /// <returns>The generated C# source code as a string.</returns>
    public static string Generate(ExecutorInfo info)
    {
        var sb = new StringBuilder();

        // File header
        sb.AppendLine("// <auto-generated/>");
        sb.AppendLine("#nullable enable");
        sb.AppendLine();

        // Using directives
        sb.AppendLine("using System;");
        sb.AppendLine("using System.Collections.Generic;");
        sb.AppendLine("using Microsoft.Agents.AI.Workflows;");
        sb.AppendLine();
        sb.AppendLine("using RouteBuilder = Microsoft.Agents.AI.Workflows.RouteBuilder;");

        // Namespace
        if (!string.IsNullOrWhiteSpace(info.Namespace))
        {
            sb.AppendLine($"namespace {info.Namespace};");
            sb.AppendLine();
        }

        // For nested classes, we must emit partial declarations for each containing type.
        // Example: if MyExecutor is nested in Outer.Inner, we emit:
        //   partial class Outer { partial class Inner { partial class MyExecutor { ... } } }
        string indent = "";
        if (info.IsNested)
        {
            foreach (string containingType in info.ContainingTypeChain.Split('.'))
            {
                sb.AppendLine($"{indent}partial class {containingType}");
                sb.AppendLine($"{indent}{{");

                indent += IndentUnit;
            }
        }

        // Class declaration
        sb.AppendLine($"{indent}partial class {info.ClassName}{info.GenericParameters}");
        sb.AppendLine($"{indent}{{");

        string memberIndent = indent + IndentUnit;

        // ConfigureProtocol
        sb.AppendLine($"{memberIndent}protected override ProtocolBuilder ConfigureProtocol(ProtocolBuilder protocolBuilder)");
        sb.AppendLine($"{memberIndent}{{");

        string bodyIndent = memberIndent + IndentUnit;

        if (info.BaseHasConfigureProtocol)
        {
            sb.Append($"{bodyIndent}return base.ConfigureProtocol(protocolBuilder)");
            bodyIndent += "           ";
        }
        else
        {
            sb.Append($"{bodyIndent}return protocolBuilder");
        }

        // Only generate protocol overrides if [SendsMessage] or [YieldsOutput] attributes are present.
        // Without these attributes, we rely on the base class defaults.
        if (info.ShouldGenerateSentMessageRegistrations)
        {
            GenerateConfigureSentTypes(sb, info, bodyIndent);
        }

        if (info.ShouldGenerateYieldedOutputRegistrations)
        {
            GenerateConfigureYieldTypes(sb, info, bodyIndent);
        }

        // Only generate ConfigureRoutes if there are handlers
        if (info.Handlers.Count > 0)
        {
            GenerateConfigureRoutes(sb, info, bodyIndent);
        }
        else
        {
            sb.AppendLine(";");
        }

        // Close ConfigureProtocol
        sb.AppendLine($"{memberIndent}}}");

        // Close class
        sb.AppendLine($"{indent}}}");

        // Close nested classes
        if (info.IsNested)
        {
            string[] containingTypes = info.ContainingTypeChain.Split('.');
            for (int i = containingTypes.Length - 1; i >= 0; i--)
            {
                indent = new string(' ', i * 4);
                sb.AppendLine($"{indent}}}");
            }
        }

        return sb.ToString();
    }

    /// <summary>
    /// Generates the ConfigureRoutes override that registers all [MessageHandler] methods.
    /// </summary>
    private static void GenerateConfigureRoutes(StringBuilder sb, ExecutorInfo info, string indent)
    {
        sb.AppendLine(".ConfigureRoutes(ConfigureRoutes);");

        sb.AppendLine($"{indent}void ConfigureRoutes(RouteBuilder routeBuilder)");
        sb.AppendLine($"{indent}{{");

        string bodyIndent = indent + IndentUnit;

        // Generate handler registrations using fluent AddHandler calls.
        // RouteBuilder.AddHandler<TIn> registers a void handler; AddHandler<TIn, TOut> registers one with a return value.
        if (info.Handlers.Count == 1)
        {
            HandlerInfo handler = info.Handlers[0];
            sb.AppendLine($"{bodyIndent}routeBuilder");
            sb.Append($"{bodyIndent}    .AddHandler");
            AppendHandlerGenericArgs(sb, handler);
            sb.AppendLine($"(this.{handler.MethodName});");
        }
        else
        {
            // Multiple handlers: chain fluent calls, semicolon only on the last one.
            sb.AppendLine($"{bodyIndent}routeBuilder");

            for (int i = 0; i < info.Handlers.Count; i++)
            {
                HandlerInfo handler = info.Handlers[i];

                sb.Append($"{bodyIndent}    .AddHandler");
                AppendHandlerGenericArgs(sb, handler);
                sb.Append($"(this.{handler.MethodName})");
                sb.AppendLine();
            }

            // Remove last newline without using that System.Environment which is banned from use in analyzers
            var newLineLength = new StringBuilder().AppendLine().Length;
            sb.Remove(sb.Length - newLineLength, newLineLength);
            sb.AppendLine(";");
        }

        sb.AppendLine($"{indent}}}");
    }

    /// <summary>
    /// Appends generic type arguments for AddHandler based on whether the handler returns a value.
    /// </summary>
    private static void AppendHandlerGenericArgs(StringBuilder sb, HandlerInfo handler)
    {
        // Handlers returning ValueTask use single type arg; ValueTask<T> uses two.
        if (handler.HasOutput && handler.OutputTypeName != null)
        {
            sb.Append($"<{handler.InputTypeName}, {handler.OutputTypeName}>");
        }
        else
        {
            sb.Append($"<{handler.InputTypeName}>");
        }
    }

    /// <summary>
    /// Generates ConfigureSentTypes override declaring message types this executor sends via context.SendMessageAsync.
    /// </summary>
    /// <remarks>
    /// Types come from [SendsMessage] attributes on the class or individual handler methods.
    /// This enables workflow protocol validation at build time.
    /// </remarks>
    private static void GenerateConfigureSentTypes(StringBuilder sb, ExecutorInfo info, string indent)
    {
        // Track types to avoid emitting duplicate Add calls (the set handles runtime dedup,
        // but cleaner generated code is easier to read).
        var addedTypes = new HashSet<string>();

        foreach (var type in info.ClassSendTypes.Where(type => addedTypes.Add(type)))
        {
            sb.AppendLine($".SendsMessage<{type}>()");
            sb.Append(indent);
        }

        foreach (var handler in info.Handlers)
        {
            foreach (var type in handler.SendTypes.Where(type => addedTypes.Add(type)))
            {
                sb.AppendLine($".SendsMessage<{type}>()");
                sb.Append(indent);
            }
        }
    }

    /// <summary>
    /// Generates ConfigureYieldTypes override declaring message types this executor yields via context.YieldOutputAsync.
    /// </summary>
    /// <remarks>
    /// Types come from [YieldsOutput] attributes and handler return types (ValueTask&lt;T&gt;).
    /// This enables workflow protocol validation at build time.
    /// </remarks>
    private static void GenerateConfigureYieldTypes(StringBuilder sb, ExecutorInfo info, string indent)
    {
        // Track types to avoid emitting duplicate Add calls (the set handles runtime dedup,
        // but cleaner generated code is easier to read).
        var addedTypes = new HashSet<string>();

        foreach (var type in info.ClassYieldTypes.Where(type => addedTypes.Add(type)))
        {
            sb.AppendLine($".YieldsOutput<{type}>()");
            sb.Append(indent);
        }

        foreach (var handler in info.Handlers)
        {
            foreach (var type in handler.YieldTypes.Where(type => addedTypes.Add(type)))
            {
                sb.AppendLine($".YieldsOutput<{type}>()");
                sb.Append(indent);
            }
        }
    }
}