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

using System.Collections.Concurrent;
using System.Collections.Immutable;
using System.Diagnostics;
using System.Reflection;
using System.Runtime.Loader;
using System.Security.Cryptography;
using System.Text;
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="AssemblyLoadContext"/> so the loaded assembly
/// and all its types can be garbage-collected after the call returns. The
/// compiled <em>image bytes</em> are cached by source hash, so a repeated task
/// skips Roslyn and only re-loads the image into a fresh collectible ALC.
///
/// The compile + invoke runs off the caller thread under
/// <c>Workflows:CSharpTimeoutSeconds</c> (default <c>TaskTimeoutSeconds</c>) and
/// observes <c>ct</c>: an async entry is awaited with cancellation, and the
/// caller always returns by the deadline. A synchronous entry that ignores the
/// token cannot be force-killed in-process; the caller stops waiting at the
/// deadline and the abandoned thread is left to finish (documented limitation).
///
/// 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);

    /// <summary>Compiled PE images keyed by source hash, so a hot task never re-runs Roslyn.</summary>
    private static readonly ConcurrentDictionary<string, byte[]> AssemblyCache = new(StringComparer.Ordinal);
    private static readonly ConcurrentQueue<string> CacheOrder = new();
    private const int MaxCachedAssemblies = 64;

    // 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();

    private readonly TimeSpan _timeout;

    public CSharpScriptExecutor(IConfiguration? config = null)
    {
        var seconds = ParseInt(config?["Workflows:CSharpTimeoutSeconds"], 0);
        if (seconds <= 0)
            seconds = ParseInt(config?["Workflows:TaskTimeoutSeconds"], 60);
        _timeout = TimeSpan.FromSeconds(Math.Max(1, seconds));
    }

    public string Language => "csharp";

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

    public async Task<ExecutionResult> ExecuteAsync(TaskInvocation invocation, CancellationToken ct)
    {
        var sw = Stopwatch.StartNew();
        if (!ExecutionHelpers.TryResolveEntryPath(invocation, out var entryPath, out var pathError))
            return Fail(pathError!, sw.Elapsed);
        if (!File.Exists(entryPath))
            return Fail($"entry file not found: {invocation.EntryFile}", sw.Elapsed);

        string source;
        try
        {
            source = File.ReadAllText(entryPath);
        }
        catch (Exception ex)
        {
            return Fail($"cannot read entry file: {ex.Message}", sw.Elapsed);
        }

        // Linked token: fires on caller cancellation OR the hard deadline. Used
        // inside the work for async entries; WaitAsync below guarantees the caller
        // returns at the deadline even if a sync entry ignores the token.
        var linked = CancellationTokenSource.CreateLinkedTokenSource(ct);
        linked.CancelAfter(_timeout);
        var token = linked.Token;

        try
        {
            var work = Task.Run(() => CompileAndInvokeAsync(source, entryPath, invocation, token), token);
            var output = await work.WaitAsync(_timeout, ct);

            sw.Stop();
            return new ExecutionResult(true, 0, string.Empty, string.Empty, output, null, sw.Elapsed);
        }
        catch (TimeoutException)
        {
            return Fail($"C# script timed out after {_timeout.TotalSeconds:0}s", sw.Elapsed);
        }
        catch (OperationCanceledException) when (ct.IsCancellationRequested)
        {
            // P1-11: cancellation is not a normal script failure.
            throw;
        }
        catch (OperationCanceledException)
        {
            return Fail($"C# script timed out after {_timeout.TotalSeconds:0}s", sw.Elapsed);
        }
        catch (Exception ex)
        {
            return Fail(ex.Message, sw.Elapsed);
        }
    }

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

    private static async Task<string?> CompileAndInvokeAsync(
        string source, string entryPath, TaskInvocation invocation, CancellationToken ct)
    {
        ct.ThrowIfCancellationRequested();
        var image = GetOrCompile(source, entryPath, ct);

        var alc = new AssemblyLoadContext(name: null, isCollectible: true);
        try
        {
            using var ms = new MemoryStream(image);
            var assembly = alc.LoadFromStream(ms);
            return await InvokeEntryAsync(assembly, invocation, ct);
        }
        finally
        {
            // Unload the collectible ALC so the assembly and all its types become
            // eligible for GC. The next GC.Collect will reclaim them.
            alc.Unload();
        }
    }

    /// <summary>
    /// Returns the compiled PE image for <paramref name="source"/>, compiling on a
    /// cache miss. Throws <see cref="CSharpCompileException"/> with the Roslyn
    /// diagnostics when compilation fails (failures are never cached).
    /// </summary>
    private static byte[] GetOrCompile(string source, string path, CancellationToken ct)
    {
        var key = Convert.ToHexString(SHA256.HashData(Encoding.UTF8.GetBytes(source)));
        if (AssemblyCache.TryGetValue(key, out var cached))
            return cached;

        var image = Compile(source, path);
        ct.ThrowIfCancellationRequested();

        if (AssemblyCache.TryAdd(key, image))
        {
            CacheOrder.Enqueue(key);
            while (CacheOrder.Count > MaxCachedAssemblies && CacheOrder.TryDequeue(out var evicted))
                AssemblyCache.TryRemove(evicted, out _);
        }

        return image;
    }

    private static byte[] Compile(string source, string path)
    {
        var syntaxTree = CSharpSyntaxTree.ParseText(source, path: 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 errors = emit.Diagnostics
                .Where(d => d.Severity == DiagnosticSeverity.Error)
                .Select(d => d.ToString());
            throw new CSharpCompileException("C# compilation failed:\n" + string.Join('\n', errors));
        }

        return ms.ToArray();
    }

    /// <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.
    /// </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 async Task<string?> InvokeEntryAsync(
        Assembly assembly, TaskInvocation invocation, CancellationToken ct)
    {
        var method = FindEntryMethod(assembly, invocation);
        if (method == null)
        {
            var function = string.IsNullOrWhiteSpace(invocation.EntryFunction) ? "Main" : invocation.EntryFunction;
            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;

        ct.ThrowIfCancellationRequested();
        var result = method.Invoke(null, arguments);

        // Await Task / Task<T> returns so async entries produce a concrete value
        // and a cancellation is observed instead of blocking the caller.
        if (result is Task task)
        {
            await task.WaitAsync(ct);
            result = task.GetType().IsGenericType
                ? task.GetType().GetProperty("Result")!.GetValue(task)
                : null;
        }

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

    private static MethodInfo? FindEntryMethod(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;
        }

        return method;
    }

    private static int ParseInt(string? raw, int fallback)
        => int.TryParse(raw, out var value) ? value : fallback;

    /// <summary>Raised on Roslyn emit failure; the message carries the diagnostics.</summary>
    private sealed class CSharpCompileException(string message) : Exception(message);
}