File: ToolBlockEmitter.cs
Web Access
Project: src\aspnetcore\src\Components\AI\gen\Microsoft.AspNetCore.Components.AI.SourceGenerators.csproj (Microsoft.AspNetCore.Components.AI.SourceGenerators)
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
 
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Text;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;
 
namespace Microsoft.AspNetCore.Components.AI.SourceGenerators;
 
internal static class ToolBlockEmitter
{
    public static void EmitHandler(SourceProductionContext spc, ToolBlockCandidate candidate)
    {
        var blockType = candidate.BlockTypeGlobal;
        var sb = new StringBuilder();
        sb.AppendLine("// <auto-generated/>");
        sb.AppendLine("#nullable enable");
        sb.AppendLine();
 
        if (!string.IsNullOrEmpty(candidate.Namespace))
        {
            sb.AppendLine($"namespace {candidate.Namespace};");
            sb.AppendLine();
        }
 
        var className = candidate.ClassName + "Handler";
 
        sb.AppendLine($"[global::System.ComponentModel.EditorBrowsable(global::System.ComponentModel.EditorBrowsableState.Never)]");
        sb.AppendLine($"internal sealed class {className}");
        sb.AppendLine($"    : global::Microsoft.AspNetCore.Components.AI.ContentBlockHandler<{blockType}>");
        sb.AppendLine("{");
        sb.AppendLine($"    public override global::Microsoft.AspNetCore.Components.AI.BlockMappingResult<{blockType}> Handle(");
        sb.AppendLine($"        global::Microsoft.AspNetCore.Components.AI.BlockMappingContext context,");
        sb.AppendLine($"        {blockType} state)");
        sb.AppendLine("    {");
        sb.AppendLine("        if (state.Result is not null)");
        sb.AppendLine("        {");
        sb.AppendLine($"            return global::Microsoft.AspNetCore.Components.AI");
        sb.AppendLine($"                .BlockMappingResult<{blockType}>.Complete();");
        sb.AppendLine("        }");
        sb.AppendLine();
        sb.AppendLine("        var shouldEmit = false;");
        sb.AppendLine();
 
        // Phase 1: Match FunctionCallContent by tool name
        sb.AppendLine("        if (state.Call is null)");
        sb.AppendLine("        {");
        sb.AppendLine("            global::Microsoft.Extensions.AI.FunctionCallContent? callContent = null;");
        sb.AppendLine("            foreach (var content in context.UnhandledContents)");
        sb.AppendLine("            {");
        sb.AppendLine("                if (content is global::Microsoft.Extensions.AI.FunctionCallContent fc");
        sb.AppendLine($"                    && fc.Name == \"{EscapeString(candidate.ToolName)}\")");
        sb.AppendLine("                {");
        sb.AppendLine("                    callContent = fc;");
        sb.AppendLine("                    break;");
        sb.AppendLine("                }");
        sb.AppendLine("            }");
        sb.AppendLine();
        sb.AppendLine("            if (callContent is not null)");
        sb.AppendLine("            {");
        sb.AppendLine("                context.MarkHandled(callContent);");
        sb.AppendLine("                state.Call = callContent;");
        sb.AppendLine("                shouldEmit = true;");
 
        // Deserialize arguments
        if (candidate.Parameters.Count > 0)
        {
            sb.AppendLine();
            sb.AppendLine("                if (callContent.Arguments is { } args)");
            sb.AppendLine("                {");
 
            foreach (var param in candidate.Parameters)
            {
                var varName = LocalName(param.PropertyName);
                var member = "state." + EscapeIdentifier(param.PropertyName);
                sb.AppendLine($"                    if (args.TryGetValue(\"{EscapeString(param.ArgumentKey)}\", out var {varName}) && {varName} is not null)");
                sb.AppendLine("                    {");
                sb.AppendLine($"                        {member} = {varName} switch");
                sb.AppendLine("                        {");
                EmitDeserialization(sb, param, varName);
                sb.AppendLine("                        };");
                sb.AppendLine("                    }");
            }
 
            sb.AppendLine("                }");
        }
 
        sb.AppendLine("            }");
        sb.AppendLine("        }");
        sb.AppendLine();
 
        // Phase 2: Match FunctionResultContent by CallId
        sb.AppendLine("        global::Microsoft.Extensions.AI.FunctionResultContent? resultContent = null;");
        sb.AppendLine("        foreach (var content in context.UnhandledContents)");
        sb.AppendLine("        {");
        sb.AppendLine("            if (content is global::Microsoft.Extensions.AI.FunctionResultContent frc");
        sb.AppendLine("                && state.Call is not null");
        sb.AppendLine("                && frc.CallId == state.Call.CallId)");
        sb.AppendLine("            {");
        sb.AppendLine("                resultContent = frc;");
        sb.AppendLine("                break;");
        sb.AppendLine("            }");
        sb.AppendLine("        }");
        sb.AppendLine();
        sb.AppendLine("        if (resultContent is not null)");
        sb.AppendLine("        {");
        sb.AppendLine("            context.MarkHandled(resultContent);");
        sb.AppendLine("            state.Result = resultContent;");
 
        // Deserialize [ToolResult] properties
        if (candidate.ResultProperties.Count > 0)
        {
            sb.AppendLine();
            sb.AppendLine("            if (resultContent.Result is not null)");
            sb.AppendLine("            {");
 
            if (candidate.ResultProperties.Count == 1)
            {
                // Single result property: map the entire Result value directly
                var rp = candidate.ResultProperties[0];
                var varName = "__result";
                var member = "state." + EscapeIdentifier(rp.PropertyName);
                sb.AppendLine($"                var {varName} = resultContent.Result;");
                sb.AppendLine($"                {member} = {varName} switch");
                sb.AppendLine("                {");
                EmitResultDeserialization(sb, rp, varName);
                sb.AppendLine("                };");
            }
            else
            {
                // Multiple result properties: treat Result as a JSON object
                sb.AppendLine("                var __resultObj = resultContent.Result switch");
                sb.AppendLine("                {");
                sb.AppendLine("                    global::System.Text.Json.JsonElement __element => __element,");
                sb.AppendLine("                    string __json => global::System.Text.Json.JsonSerializer.Deserialize<global::System.Text.Json.JsonElement>(__json),");
                sb.AppendLine("                    _ => global::System.Text.Json.JsonSerializer.SerializeToElement(resultContent.Result),");
                sb.AppendLine("                };");
                sb.AppendLine();
                sb.AppendLine("                if (__resultObj.ValueKind == global::System.Text.Json.JsonValueKind.Object)");
                sb.AppendLine("                {");
 
                foreach (var rp in candidate.ResultProperties)
                {
                    var varName = "__r_" + LocalSuffix(rp.PropertyName);
                    var member = "state." + EscapeIdentifier(rp.PropertyName);
                    sb.AppendLine($"                    if (__resultObj.TryGetProperty(\"{EscapeString(rp.ResultKey)}\", out var {varName}))");
                    sb.AppendLine("                    {");
                    sb.AppendLine($"                        {member} = {varName} switch");
                    sb.AppendLine("                        {");
                    EmitJsonElementDeserialization(sb, rp, varName);
                    sb.AppendLine("                        };");
                    sb.AppendLine("                    }");
                }
 
                sb.AppendLine("                }");
            }
 
            sb.AppendLine("            }");
        }
 
        sb.AppendLine("            return shouldEmit");
        sb.AppendLine($"                ? global::Microsoft.AspNetCore.Components.AI.BlockMappingResult<{blockType}>.Emit(state, state)");
        sb.AppendLine($"                : global::Microsoft.AspNetCore.Components.AI.BlockMappingResult<{blockType}>.Complete();");
        sb.AppendLine("        }");
        sb.AppendLine();
        sb.AppendLine("        return shouldEmit");
        sb.AppendLine($"            ? global::Microsoft.AspNetCore.Components.AI.BlockMappingResult<{blockType}>.Emit(state, state)");
        sb.AppendLine($"            : global::Microsoft.AspNetCore.Components.AI.BlockMappingResult<{blockType}>.Pass();");
        sb.AppendLine("    }");
        sb.AppendLine("}");
 
        spc.AddSource(HintName(candidate), sb.ToString());
    }
 
    public static void EmitRegistration(SourceProductionContext spc, ImmutableArray<ToolBlockCandidate> candidates)
    {
        if (candidates.IsEmpty)
        {
            return;
        }
 
        // Check for duplicate tool names
        var seen = new Dictionary<string, ToolBlockCandidate>();
        foreach (var c in candidates)
        {
            if (seen.TryGetValue(c.ToolName, out var existing))
            {
                spc.ReportDiagnostic(Diagnostic.Create(
                    DiagnosticDescriptors.DuplicateToolName,
                    Location.None,
                    c.ToolName, existing.ClassName, c.ClassName));
            }
            else
            {
                seen[c.ToolName] = c;
            }
        }
 
        var sb = new StringBuilder();
        sb.AppendLine("// <auto-generated/>");
        sb.AppendLine("#nullable enable");
        sb.AppendLine();
        sb.AppendLine("namespace Microsoft.AspNetCore.Components.AI;");
        sb.AppendLine();
        sb.AppendLine("[global::System.ComponentModel.EditorBrowsable(global::System.ComponentModel.EditorBrowsableState.Never)]");
        sb.AppendLine("internal static class GeneratedToolBlockRegistrations");
        sb.AppendLine("{");
        sb.AppendLine("    internal static void AddGeneratedToolBlocks(this UIAgentOptions options)");
        sb.AppendLine("    {");
 
        foreach (var candidate in candidates)
        {
            var handlerName = candidate.ClassName + "Handler";
            var fullClass = string.IsNullOrEmpty(candidate.Namespace)
                ? "global::" + handlerName
                : "global::" + candidate.Namespace + "." + handlerName;
            sb.AppendLine($"        options.AddBlockHandler(new {fullClass}());");
        }
 
        sb.AppendLine("    }");
        sb.AppendLine("}");
 
        spc.AddSource("GeneratedToolBlockRegistrations.g.cs", sb.ToString());
    }
 
    // Hint names must be unique per generated file; two ToolBlocks with the same simple name in
    // different namespaces would otherwise collide. Qualify with the (sanitized) namespace while
    // keeping "<ClassName>Handler.g.cs" as the suffix.
    private static string HintName(ToolBlockCandidate candidate)
    {
        var suffix = candidate.ClassName + "Handler.g.cs";
        return string.IsNullOrEmpty(candidate.Namespace)
            ? suffix
            : Sanitize(candidate.Namespace) + "." + suffix;
    }
 
    private static string Sanitize(string value)
    {
        var sb = new StringBuilder(value.Length);
        foreach (var ch in value)
        {
            sb.Append(char.IsLetterOrDigit(ch) || ch == '.' || ch == '_' ? ch : '_');
        }
        return sb.ToString();
    }
 
    private static string EscapeIdentifier(string name)
        => SyntaxFacts.GetKeywordKind(name) != SyntaxKind.None
            || SyntaxFacts.GetContextualKeywordKind(name) != SyntaxKind.None
            ? "@" + name
            : name;
 
    private static string LocalName(string propertyName)
        => "__" + LocalSuffix(propertyName);
 
    private static string LocalSuffix(string propertyName)
        => propertyName.Substring(0, 1).ToLowerInvariant() + propertyName.Substring(1);
 
    private static void EmitDeserialization(StringBuilder sb, ToolParameterInfo param, string varName)
    {
        switch (param.TypeKind)
        {
            case ParameterTypeKind.String:
                sb.AppendLine($"                                global::System.Text.Json.JsonElement __je => __je.GetString()!,");
                sb.AppendLine($"                                string __s => __s,");
                sb.AppendLine($"                                _ => (string){varName}!");
                break;
            case ParameterTypeKind.Int32:
                sb.AppendLine($"                                global::System.Text.Json.JsonElement __je => __je.GetInt32(),");
                sb.AppendLine($"                                _ => global::System.Convert.ToInt32({varName})");
                break;
            case ParameterTypeKind.Int64:
                sb.AppendLine($"                                global::System.Text.Json.JsonElement __je => __je.GetInt64(),");
                sb.AppendLine($"                                _ => global::System.Convert.ToInt64({varName})");
                break;
            case ParameterTypeKind.Double:
                sb.AppendLine($"                                global::System.Text.Json.JsonElement __je => __je.GetDouble(),");
                sb.AppendLine($"                                _ => global::System.Convert.ToDouble({varName})");
                break;
            case ParameterTypeKind.Single:
                sb.AppendLine($"                                global::System.Text.Json.JsonElement __je => __je.GetSingle(),");
                sb.AppendLine($"                                _ => global::System.Convert.ToSingle({varName})");
                break;
            case ParameterTypeKind.Decimal:
                sb.AppendLine($"                                global::System.Text.Json.JsonElement __je => __je.GetDecimal(),");
                sb.AppendLine($"                                _ => global::System.Convert.ToDecimal({varName})");
                break;
            case ParameterTypeKind.Boolean:
                sb.AppendLine($"                                global::System.Text.Json.JsonElement __je => __je.GetBoolean(),");
                sb.AppendLine($"                                _ => global::System.Convert.ToBoolean({varName})");
                break;
            case ParameterTypeKind.Complex:
                sb.AppendLine($"                                global::System.Text.Json.JsonElement __je => global::System.Text.Json.JsonSerializer.Deserialize<{param.TypeName}>(__je)!,");
                sb.AppendLine($"                                _ => ({param.TypeName}){varName}!");
                break;
        }
    }
 
    private static string EscapeString(string value)
    {
        return value.Replace("\\", "\\\\").Replace("\"", "\\\"");
    }
 
    private static void EmitResultDeserialization(StringBuilder sb, ToolResultPropertyInfo prop, string varName)
    {
        // Single-property result: map the entire Result value directly
        switch (prop.TypeKind)
        {
            case ParameterTypeKind.String:
                sb.AppendLine($"                        global::System.Text.Json.JsonElement __je => __je.GetString()!,");
                sb.AppendLine($"                        string __s => __s,");
                sb.AppendLine($"                        _ => {varName}!.ToString()!");
                break;
            case ParameterTypeKind.Int32:
                sb.AppendLine($"                        global::System.Text.Json.JsonElement __je => __je.GetInt32(),");
                sb.AppendLine($"                        _ => global::System.Convert.ToInt32({varName})");
                break;
            case ParameterTypeKind.Int64:
                sb.AppendLine($"                        global::System.Text.Json.JsonElement __je => __je.GetInt64(),");
                sb.AppendLine($"                        _ => global::System.Convert.ToInt64({varName})");
                break;
            case ParameterTypeKind.Double:
                sb.AppendLine($"                        global::System.Text.Json.JsonElement __je => __je.GetDouble(),");
                sb.AppendLine($"                        _ => global::System.Convert.ToDouble({varName})");
                break;
            case ParameterTypeKind.Single:
                sb.AppendLine($"                        global::System.Text.Json.JsonElement __je => __je.GetSingle(),");
                sb.AppendLine($"                        _ => global::System.Convert.ToSingle({varName})");
                break;
            case ParameterTypeKind.Decimal:
                sb.AppendLine($"                        global::System.Text.Json.JsonElement __je => __je.GetDecimal(),");
                sb.AppendLine($"                        _ => global::System.Convert.ToDecimal({varName})");
                break;
            case ParameterTypeKind.Boolean:
                sb.AppendLine($"                        global::System.Text.Json.JsonElement __je => __je.GetBoolean(),");
                sb.AppendLine($"                        _ => global::System.Convert.ToBoolean({varName})");
                break;
            case ParameterTypeKind.Complex:
                sb.AppendLine($"                        global::System.Text.Json.JsonElement __je => global::System.Text.Json.JsonSerializer.Deserialize<{prop.TypeName}>(__je)!,");
                sb.AppendLine($"                        string __json => global::System.Text.Json.JsonSerializer.Deserialize<{prop.TypeName}>(__json)!,");
                sb.AppendLine($"                        _ => ({prop.TypeName}){varName}!");
                break;
        }
    }
 
    private static void EmitJsonElementDeserialization(StringBuilder sb, ToolResultPropertyInfo prop, string varName)
    {
        // Multi-property result: deserialize from a JSON object property
        switch (prop.TypeKind)
        {
            case ParameterTypeKind.String:
                sb.AppendLine($"                                global::System.Text.Json.JsonElement {{ ValueKind: global::System.Text.Json.JsonValueKind.String }} => {varName}.GetString()!,");
                sb.AppendLine($"                                _ => {varName}.GetRawText()");
                break;
            case ParameterTypeKind.Int32:
                sb.AppendLine($"                                _ => {varName}.GetInt32()");
                break;
            case ParameterTypeKind.Int64:
                sb.AppendLine($"                                _ => {varName}.GetInt64()");
                break;
            case ParameterTypeKind.Double:
                sb.AppendLine($"                                _ => {varName}.GetDouble()");
                break;
            case ParameterTypeKind.Single:
                sb.AppendLine($"                                _ => {varName}.GetSingle()");
                break;
            case ParameterTypeKind.Decimal:
                sb.AppendLine($"                                _ => {varName}.GetDecimal()");
                break;
            case ParameterTypeKind.Boolean:
                sb.AppendLine($"                                _ => {varName}.GetBoolean()");
                break;
            case ParameterTypeKind.Complex:
                sb.AppendLine($"                                _ => global::System.Text.Json.JsonSerializer.Deserialize<{prop.TypeName}>({varName})!");
                break;
        }
    }
}