w4c-workflows-api/Services/Execution/CSharpScriptExecutor.cs

using System.Collections.Immutable;
using System.Diagnostics;
using System.Reflection;
using System.Runtime.Loader;
using System.Text.Json;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;

namespace w4c_workflows.Services.Execution;

/// <summary>
/// Executes single-file C# with the Roslyn compiler: the entry <c>.cs</c> file
/// specified by the task's <c>entry.file</c> is compiled into an in-memory
/// assembly, then the entry method is invoked reflectively.
///
/// Only the entry file is compiled — not every <c>.cs</c> file under the
/// working directory. Compiling all files would merge unrelated task sources
/// (e.g. <c>Main.cs</c> + <c>tasks/hello-world.cs</c>) into one assembly,
/// causing duplicate-type errors when both define a <c>Program</c> class.
///
/// Each execution uses a <see cref="CollectibleAssemblyLoadContext"/> so the
/// loaded assembly and all its types can be garbage-collected after the call
/// returns. The previous implementation loaded into the default (non-collectible)
/// ALC, causing unbounded memory growth in long-running workers.
///
/// Entry contract (v1): a <c>public static</c> method named after
/// <c>entry.function</c> (default <c>Main</c>) that takes a single
/// <c>string</c> or <c>object</c>/<c>dynamic</c> argument (the JSON input) or
/// no arguments, returning a value serialized to JSON as the task output (sync
/// or <c>Task</c>/<c>Task&lt;T&gt;</c>).
/// </summary>
public class CSharpScriptExecutor : IScriptExecutor
{
    private static readonly JsonSerializerOptions Json = new(JsonSerializerDefaults.Web);

    // Cached once per process lifetime — the set of runtime assemblies doesn't
    // change after startup. This avoids re-scanning AppDomain on every compile.
    private static ImmutableArray<MetadataReference>? _cachedReferences;
    private static readonly object _referenceLock = new();

    public string Language => "csharp";

    /// <summary>Roslyn ships with the service image — no external runtime to probe.</summary>
    public bool IsAvailable() => true;

    public Task<ExecutionResult> ExecuteAsync(TaskInvocation invocation, CancellationToken ct)
    {
        var sw = Stopwatch.StartNew();
        try
        {
            var workingDir = ExecutionHelpers.ResolveWorkingDir(invocation);
            var entryPath = Path.GetFullPath(Path.Combine(workingDir, invocation.EntryFile));
            if (!File.Exists(entryPath))
                return Task.FromResult(Fail($"entry file not found: {invocation.EntryFile}", sw.Elapsed));

            var source = new SourceFile(invocation.EntryFile, File.ReadAllText(entryPath));

            // Compilation is CPU-bound and synchronous; keep it off the request thread.
            var (assemblyLoadContext, assembly, diagnostics) = Compile(source);
            if (assembly == null)
            {
                assemblyLoadContext?.Unload();
                var errors = diagnostics
                    .Where(d => d.Severity == DiagnosticSeverity.Error)
                    .Select(d => d.ToString());
                return Task.FromResult(Fail("C# compilation failed:\n" + string.Join('\n', errors), sw.Elapsed));
            }

            string? output;
            try
            {
                output = InvokeEntry(assembly, invocation);
            }
            finally
            {
                // Unload the collectible ALC so the assembly and all its types
                // become eligible for GC. The next GC.Collect will reclaim them.
                assemblyLoadContext.Unload();
            }

            sw.Stop();
            return Task.FromResult(new ExecutionResult(true, 0, string.Empty, string.Empty, output, null, sw.Elapsed));
        }
        catch (Exception ex)
        {
            sw.Stop();
            return Task.FromResult(new ExecutionResult(false, -1, string.Empty, string.Empty, null, ex.Message, sw.Elapsed));
        }
    }

    private static ExecutionResult Fail(string message, TimeSpan duration)
        => new(false, -1, string.Empty, string.Empty, null, message, duration);

    private static (AssemblyLoadContext, Assembly?, ImmutableArray<Diagnostic>) Compile(SourceFile source)
    {
        var syntaxTree = CSharpSyntaxTree.ParseText(source.Source, path: source.Path);

        var compilation = CSharpCompilation.Create(
            $"wf-cs-{Guid.NewGuid():N}",
            [syntaxTree],
            GetCachedReferences(),
            new CSharpCompilationOptions(OutputKind.DynamicallyLinkedLibrary, allowUnsafe: true));

        using var ms = new MemoryStream();
        var emit = compilation.Emit(ms);
        if (!emit.Success)
        {
            var emptyAlc = new AssemblyLoadContext(name: null, isCollectible: true);
            return (emptyAlc, null, emit.Diagnostics);
        }

        ms.Position = 0;

        // Each execution gets its own collectible ALC. After Unload() + GC,
        // the assembly and all its types are reclaimed — no unbounded growth.
        var alc = new AssemblyLoadContext(name: null, isCollectible: true);
        var assembly = alc.LoadFromStream(ms);
        return (alc, assembly, emit.Diagnostics);
    }

    /// <summary>
    /// Returns the cached set of metadata references for the current runtime.
    /// The first call scans AppDomain; subsequent calls return the cached
    /// immutable array. This eliminates the per-compile AppDomain scan that
    /// compounded the ALC leak.
    /// </summary>
    private static ImmutableArray<MetadataReference> GetCachedReferences()
    {
        if (_cachedReferences is { } cached)
            return cached;

        lock (_referenceLock)
        {
            if (_cachedReferences is { } cached2)
                return cached2;

            // Assemblies injected by dotnet-watch hot reload (Edit and Continue)
            // contain duplicate metadata keys that crash Roslyn's internal cache.
            // Skip them — user code never references these types directly.
            static bool IsHotReloadAssembly(Assembly asm)
            {
                var name = asm.GetName().Name;
                if (string.IsNullOrEmpty(name))
                    return false;
                return name.StartsWith("Microsoft.CodeAnalysis", StringComparison.Ordinal)
                    || name.StartsWith("System.Reflection.Metadata", StringComparison.Ordinal)
                    || name.Contains("HotReload", StringComparison.OrdinalIgnoreCase)
                    || name.Contains("EditAndContinue", StringComparison.OrdinalIgnoreCase);
            }

            var references = new List<MetadataReference>();
            var seen = new HashSet<string>(StringComparer.OrdinalIgnoreCase);
            foreach (var asm in AppDomain.CurrentDomain.GetAssemblies())
            {
                if (asm.IsDynamic || IsHotReloadAssembly(asm))
                    continue;
                var location = asm.Location;
                if (string.IsNullOrEmpty(location) || !File.Exists(location) || !seen.Add(location))
                    continue;
                try
                {
                    references.Add(MetadataReference.CreateFromFile(location));
                }
                catch
                {
                    // Some assemblies (e.g. reflection-only) cannot be referenced — skip.
                }
            }

            _cachedReferences = ImmutableArray.CreateRange(references);
            return _cachedReferences.Value;
        }
    }

    private static string? InvokeEntry(Assembly assembly, TaskInvocation invocation)
    {
        var function = string.IsNullOrWhiteSpace(invocation.EntryFunction) ? "Main" : invocation.EntryFunction;

        MethodInfo? method = null;
        foreach (var type in assembly.GetTypes())
        {
            foreach (var candidate in type.GetMethods(BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Static))
            {
                if (!candidate.IsStatic || !string.Equals(candidate.Name, function, StringComparison.Ordinal))
                    continue;
                var parameters = candidate.GetParameters();
                // Accept a single JSON-input parameter typed as `string` or
                // `object`. The frontend/LLM stubs declare `dynamic input`, which
                // compiles to `object` at runtime — a `string`-only match would
                // reject every UI/AI-generated entry. Both receive the raw JSON
                // as the argument; a `dynamic`/`object` param lets the author
                // parse it or echo it back.
                var singleInput = parameters.Length == 1
                    && (parameters[0].ParameterType == typeof(string) || parameters[0].ParameterType == typeof(object));
                if (singleInput)
                {
                    method = candidate; // best match — one JSON-input parameter
                    break;
                }
                method ??= parameters.Length == 0 ? candidate : null;
            }
            if (method != null && method.GetParameters().Length == 1)
                break;
        }

        if (method == null)
            throw new InvalidOperationException($"no static entry method '{function}' found in the compiled C# program");

        var arguments = method.GetParameters().Length == 1
            ? new object?[] { invocation.Input ?? "{}" }
            : null;
        var result = method.Invoke(null, arguments);

        // Await Task / Task<T> returns so async entries produce a concrete value.
        if (result is Task task)
        {
            task.GetAwaiter().GetResult();
            result = task.GetType().IsGenericType
                ? task.GetType().GetProperty("Result")!.GetValue(task)
                : null;
        }

        return result == null ? null : JsonSerializer.Serialize(result, Json);
    }

    private sealed record SourceFile(string Path, string Source);
}