ScenarioCodeCompiler.cs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

using System.Diagnostics;
using System.Runtime.Loader;
using System.Text;
using System.Text.Json;
using System.Xml.Linq;
using Lilith.Core.Tools;

namespace Lilith.Core.Scenarios;

/// <summary>Compiles trusted scenario C# using the installed SDK. This is not a sandbox.</summary>
internal sealed class ScenarioCodeCompiler : IDisposable
{
    private AssemblyLoadContext? _context;
    public List<ILilithTool> Tools { get; } = [];

    public static async Task<ScenarioCodeCompiler> CompileAsync(ScenarioCodeTool[] definitions,
        IReadOnlyDictionary<string, string> inputs, CancellationToken cancellationToken)
    {
        var compiler = new ScenarioCodeCompiler();
        if (definitions.Length == 0) return compiler;
        try
        {
            string folder = Path.Combine(Path.GetTempPath(), "Lilith-ScenarioTools", Guid.NewGuid().ToString("N"));
            Directory.CreateDirectory(folder);
            string project = Path.Combine(folder, "ScenarioCode.csproj");
            var xml = new XDocument(new XElement("Project", new XAttribute("Sdk", "Microsoft.NET.Sdk"),
                new XElement("PropertyGroup", new XElement("TargetFramework", $"net{Environment.Version.Major}.0"),
                    new XElement("ImplicitUsings", "enable"), new XElement("Nullable", "enable"),
                    new XElement("EnableDefaultCompileItems", "false")),
                new XElement("ItemGroup", new XElement("Compile", new XAttribute("Include", "Code.cs")))));
            xml.Save(project);
            // Framework-only build: no package feeds or arbitrary project/MSBuild XML from scenarios.
            await File.WriteAllTextAsync(Path.Combine(folder, "NuGet.Config"),
                "<configuration><packageSources><clear /></packageSources></configuration>", cancellationToken);
            var code = new StringBuilder("using System.Text.Json;\npublic static class ScenarioFunctions\n{\n");
            for (int i = 0; i < definitions.Length; i++)
            {
                code.AppendLine("#pragma warning disable CS1998");
                code.AppendLine($"public static async Task<string> Tool{i}(JsonElement arguments, IReadOnlyDictionary<string,string> inputs, CancellationToken cancellationToken) {{");
                code.AppendLine(definitions[i].Code);
                code.AppendLine("}");
            }
            code.AppendLine("}");
            await File.WriteAllTextAsync(Path.Combine(folder, "Code.cs"), code.ToString(), cancellationToken);
            var start = new ProcessStartInfo("dotnet")
            {
                WorkingDirectory = folder, UseShellExecute = false, CreateNoWindow = true,
                RedirectStandardOutput = true, RedirectStandardError = true
            };
            foreach (string argument in new[] { "build", project, "-c", "Release", "--nologo", "--disable-build-servers", "-v", "quiet" })
                start.ArgumentList.Add(argument);
            using var process = Process.Start(start) ?? throw new InvalidOperationException("Could not start C# compiler.");
            Task<string> output = process.StandardOutput.ReadToEndAsync();
            Task<string> error = process.StandardError.ReadToEndAsync();
            using var deadline = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
            deadline.CancelAfter(TimeSpan.FromMinutes(2));
            try { await process.WaitForExitAsync(deadline.Token); }
            catch
            {
                try { if (!process.HasExited) process.Kill(entireProcessTree: true); } catch (InvalidOperationException) { }
                await process.WaitForExitAsync();
                await Task.WhenAll(output, error);
                throw;
            }
            string diagnostics = await output + await error;
            if (process.ExitCode != 0) throw new InvalidDataException("Scenario C# compilation failed: " + diagnostics[..Math.Min(4000, diagnostics.Length)]);
            compiler._context = new AssemblyLoadContext("Scenario-" + Guid.NewGuid().ToString("N"), isCollectible: true);
            var assembly = compiler._context.LoadFromAssemblyPath(Path.Combine(folder, "bin", "Release", $"net{Environment.Version.Major}.0", "ScenarioCode.dll"));
            var type = assembly.GetType("ScenarioFunctions", throwOnError: true)!;
            for (int i = 0; i < definitions.Length; i++)
            {
                var method = type.GetMethod("Tool" + i)!;
                var handler = method.CreateDelegate<Func<JsonElement, IReadOnlyDictionary<string, string>, CancellationToken, Task<string>>>();
                compiler.Tools.Add(new CompiledTool(definitions[i], inputs, handler));
            }
            return compiler;
        }
        catch { compiler.Dispose(); throw; }
    }

    public void Dispose()
    {
        Tools.Clear();
        _context?.Unload();
        _context = null;
    }

    private sealed class CompiledTool(ScenarioCodeTool definition, IReadOnlyDictionary<string, string> inputs,
        Func<JsonElement, IReadOnlyDictionary<string, string>, CancellationToken, Task<string>> handler) : ILilithTool
    {
        public string Name => definition.Name;
        public string Description => definition.Description;
        public JsonElement Parameters => definition.Parameters;
        public Task<string> ExecuteAsync(JsonElement arguments, CancellationToken cancellationToken = default) =>
            handler(arguments, inputs, cancellationToken);
    }
}