Genarrative/plugins/agc-unity-editor/dotnet/src/HelperHost.cs

using System.Text;
using System.Text.Json;
using System.Text.Json.Nodes;
using System.Text.RegularExpressions;
using DotCraft.Unity;

namespace Agc.Unity.Attach;

public sealed class HelperHost(IEditorBackend backend) : IAsyncDisposable
{
    public const string Protocol = "agc.unity.attach.v1";
    public const int MaximumCodeBytes = 128 * 1024;
    public const int MaximumMessageBytes = 2 * 1024 * 1024;
    private readonly SemaphoreSlim operationGate = new(1, 1);
    private Connection? connection;
    private int reconciliationRequired;
    private sealed record Connection(EditorTarget Target, string ProjectPath, string Version, string Generation);

    public async Task<JsonObject> HandleAsync(JsonObject request)
    {
        long? id = null;
        string? method = null;
        Request? input = null;
        try
        {
            id = ReadId(request);
            method = ReadString(request, "method");
            if (ReadString(request, "jsonrpc") != "2.0" || ReadString(request, "protocol") != Protocol)
                throw new HostException("UnityProtocolInvalid", "Unsupported helper protocol.");
            if (method is not ("detect" or "connect" or "status" or "execute" or "disconnect"))
                throw new HostException("UnityMethodInvalid", "Unsupported helper method.");
            if (method == "execute" && Volatile.Read(ref reconciliationRequired) != 0)
                return Reply(id, NeedsReconciliation("UnityExecutionBlocked", "An earlier execution needs reconciliation. Restart the host only after checking Unity."));
            input = ParseRequest(method, request["params"] as JsonObject);
        }
        catch (Exception error)
        {
            return method == "execute" && id != null
                ? Reply(id, Failure(error, false)) : ErrorReply(id, error);
        }

        if (!await operationGate.WaitAsync(0))
        {
            var busy = new HostException("UnityBusy", "Another Unity request is still active.");
            return method == "execute" ? Reply(id, Failure(busy, false)) : ErrorReply(id, busy);
        }

        using var deadline = new CancellationTokenSource(input.TimeoutMs);
        var dispatch = new DispatchState();
        var operation = Task.Run(async () =>
        {
            try
            {
                var value = await RunAsync(input, dispatch, deadline.Token);
                if (Encoding.UTF8.GetByteCount(value.ToJsonString()) > MaximumMessageBytes - 256)
                {
                    if (input.Method == "execute") return Latch("UnityResponseTooLarge", "The Unity result exceeds the host message limit.");
                    throw new HostException("UnityResponseTooLarge", "The Unity result exceeds the host message limit.");
                }
                return value;
            }
            catch (Exception error)
            {
                dispatch.Close();
                ClearConnection();
                if (input.Method == "execute")
                    return dispatch.Sent
                        ? Latch("UnityOutcomeUnknown", "The execution outcome is unknown; inspect Unity before any further execution.")
                        : Failure(error, false);
                throw;
            }
            finally
            {
                if (deadline.IsCancellationRequested) ClearConnection();
                operationGate.Release();
            }
        });
        try
        {
            return Reply(id, await operation.WaitAsync(deadline.Token));
        }
        catch (Exception error)
        {
            dispatch.Close();
            deadline.Cancel();
            ClearConnection();
            // An uncooperative compilation or native call retains the gate until it has stopped.
            // Its cancelled token and dispatch callback prevent any delayed execution dispatch.
            _ = operation.ContinueWith(task => _ = task.Exception, TaskContinuationOptions.OnlyOnFaulted);
            if (method == "execute")
            {
                if (dispatch.Sent) return Reply(id, Latch("UnityOutcomeUnknown", "The execution outcome is unknown; inspect Unity before any further execution."));
                return Reply(id, Failure(error, false));
            }
            return ErrorReply(id, error);
        }
    }

    private async Task<JsonObject> RunAsync(Request input, DispatchState dispatch, CancellationToken cancellationToken)
    {
        cancellationToken.ThrowIfCancellationRequested();
        if (input.Method == "connect") ClearConnection();
        if (connection != null && !ProjectIdentity.Same(connection.ProjectPath, input.ProjectPath)) ClearConnection();
        var project = ProjectIdentity.Normalize(input.ProjectPath);
        if (connection != null && !ProjectIdentity.Same(connection.ProjectPath, project)) ClearConnection();
        switch (input.Method)
        {
            case "detect":
                return Info(project, await SelectAsync(project, input.ProcessId, false, cancellationToken), null);
            case "connect":
                return Info(project, null, await ConnectAsync(project, input.ProcessId, cancellationToken));
            case "status":
                return await StatusAsync(project, input.ProcessId, cancellationToken);
            case "disconnect":
                var previous = connection;
                connection = null;
                try { if (previous != null) await backend.DisconnectAsync(cancellationToken); }
                finally { backend.ForgetSelection(); }
                return Info(project, null, null);
            case "execute":
                if (Volatile.Read(ref reconciliationRequired) != 0)
                    return NeedsReconciliation("UnityExecutionBlocked", "An earlier execution needs reconciliation.");
                await ConnectAsync(project, input.ProcessId, cancellationToken);
                cancellationToken.ThrowIfCancellationRequested();
                var reply = await backend.ExecuteAsync(input.Code!, () =>
                {
                    dispatch.MarkSent(cancellationToken);
                }, cancellationToken);
                string? executionId = null;
                while (true)
                {
                    cancellationToken.ThrowIfCancellationRequested();
                    var state = ReadOptionalString(reply, "state");
                    var generation = ReadOptionalString(reply, "generation");
                    var replyId = ReadOptionalString(reply, "executionId");
                    if (state is "unknown" or "lost")
                        return Latch("UnityOutcomeUnknown", "Unity lost the execution outcome; inspect the Editor before continuing.", replyId, generation);
                    if (!dispatch.Sent || string.IsNullOrEmpty(generation) || string.IsNullOrEmpty(replyId)
                        || (executionId != null && executionId != replyId))
                        return Latch("UnityReplyInvalid", "Unity returned an invalid execution receipt.", executionId, generation);
                    executionId = replyId;
                    switch (state)
                    {
                        case "completed":
                            if (!reply.ContainsKey("result")) return Latch("UnityReplyInvalid", "Unity returned no execution result.", executionId, generation);
                            if (connection != null) connection = connection with { Generation = generation };
                            return new JsonObject
                            {
                                ["ok"] = true, ["status"] = "completed", ["retryAllowed"] = false, ["dispatched"] = true,
                                ["executionId"] = executionId, ["generation"] = generation, ["result"] = reply["result"]?.DeepClone()
                            };
                        case "failed":
                        case "cancelled":
                            var failure = Failure(new HostException(ReadOptionalString(reply, "errorCode") ?? "UnityExecutionFailed",
                                ReadOptionalString(reply, "error") ?? "Unity did not complete the execution."), true);
                            failure["executionId"] = executionId;
                            failure["generation"] = generation;
                            return failure;
                        case "queued":
                        case "running":
                            reply = await backend.WaitAsync(executionId, cancellationToken);
                            break;
                        default:
                            return Latch("UnityReplyInvalid", "Unity returned an unsupported execution state.", executionId, generation);
                    }
                }
            default: throw new HostException("UnityMethodInvalid", "Unsupported helper method.");
        }
    }

    private async Task<Connection> ConnectAsync(string project, int? pid, CancellationToken cancellationToken)
    {
        ClearConnection();
        try
        {
            var target = (await SelectAsync(project, pid, true, cancellationToken))!;
            var metadata = await backend.ConnectAsync(target, cancellationToken);
            cancellationToken.ThrowIfCancellationRequested();
            var current = await SelectAsync(project, target.Pid, true, cancellationToken);
            if (current!.StartedUtc != target.StartedUtc)
                throw new HostException("UnityTargetChanged", "The selected Unity process changed during connection.");
            var result = VerifyMetadata(project, metadata);
            var connected = new Connection(target, project, ReadString(result, "version"), ReadString(metadata, "generation"));
            cancellationToken.ThrowIfCancellationRequested();
            connection = connected;
            return connected;
        }
        catch { ClearConnection(); throw; }
    }

    private async Task<JsonObject> StatusAsync(string project, int? pid, CancellationToken cancellationToken)
    {
        var previous = connection;
        if (previous == null) return Info(project, null, null);
        try
        {
            var current = await SelectAsync(project, pid ?? previous.Target.Pid, true, cancellationToken);
            if (current!.Pid != previous.Target.Pid || current.StartedUtc != previous.Target.StartedUtc)
                throw new HostException("UnityTargetChanged", "The connected Unity process changed.");
            var metadata = await backend.StatusAsync(cancellationToken);
            var result = VerifyMetadata(project, metadata);
            cancellationToken.ThrowIfCancellationRequested();
            connection = previous with { Version = ReadString(result, "version"), Generation = ReadString(metadata, "generation") };
            return Info(project, null, connection);
        }
        catch { ClearConnection(); throw; }
    }

    private async Task<EditorTarget?> SelectAsync(string project, int? pid, bool required, CancellationToken cancellationToken)
    {
        var targets = await backend.DetectAsync(cancellationToken);
        var candidates = targets.Where(target => ProjectIdentity.Same(target.ProjectPath, project)
            && (!pid.HasValue || target.Pid == pid)).ToArray();
        if (candidates.Length > 1) throw new HostException("UnityTargetAmbiguous", "Multiple Unity Editors match the project; select an explicit process.");
        if (candidates.Length == 0)
        {
            if (pid.HasValue || required) throw new HostException("UnityTargetUnavailable", "No Unity Editor matches the requested project and process identity.");
            return null;
        }
        if (candidates[0].Pid <= 0 || candidates[0].StartedUtc == default)
            throw new HostException("UnityTargetInvalid", "Unity process identity could not be verified.");
        return candidates[0];
    }

    private static JsonObject VerifyMetadata(string project, JsonObject reply)
    {
        if (ReadOptionalString(reply, "state") != "completed" || reply["result"] is not JsonObject result
            || !ProjectIdentity.Same(ReadOptionalString(result, "project"), project)
            || string.IsNullOrEmpty(ReadOptionalString(reply, "generation")) || string.IsNullOrEmpty(ReadOptionalString(result, "version")))
            throw new HostException("UnityHandshakeInvalid", "Unity handshake did not match the selected project.");
        return result;
    }

    private void ClearConnection() { connection = null; backend.ForgetSelection(); }
    public void MarkOutputUncertain() => Interlocked.Exchange(ref reconciliationRequired, 1);
    private JsonObject Latch(string code, string message, string? executionId = null, string? generation = null)
    {
        Interlocked.Exchange(ref reconciliationRequired, 1);
        ClearConnection();
        return NeedsReconciliation(code, message, executionId, generation);
    }

    private static JsonObject Info(string project, EditorTarget? discovered, Connection? connected) => new()
    {
        ["adapter"] = "unity-editor", ["connected"] = connected != null,
        ["pid"] = connected?.Target.Pid ?? discovered?.Pid, ["projectPath"] = project,
        ["version"] = connected?.Version, ["startedUtc"] = connected?.Target.StartedUtc ?? discovered?.StartedUtc,
        ["generation"] = connected?.Generation
    };

    private static JsonObject NeedsReconciliation(string code, string message, string? executionId = null, string? generation = null) => new()
    {
        ["ok"] = false, ["status"] = "needs-reconciliation", ["retryAllowed"] = false, ["dispatched"] = true,
        ["error"] = new JsonObject { ["code"] = code, ["message"] = message }, ["executionId"] = executionId, ["generation"] = generation
    };

    private static JsonObject Failure(Exception error, bool dispatched) => new()
    {
        ["ok"] = false, ["status"] = "failed", ["retryAllowed"] = false, ["dispatched"] = dispatched, ["error"] = Describe(error)
    };
    public static JsonObject Reply(long? id, JsonObject result) => new()
    {
        ["jsonrpc"] = "2.0", ["protocol"] = Protocol, ["id"] = id, ["result"] = result
    };
    public static JsonObject ErrorReply(long? id, Exception error) => new()
    {
        ["jsonrpc"] = "2.0", ["protocol"] = Protocol, ["id"] = id, ["error"] = Describe(error)
    };

    private static JsonObject Describe(Exception error)
    {
        var code = error switch
        {
            HostException host => host.Code, UnityTargetException target => target.Code,
            OperationCanceledException or TimeoutException => "UnityDeadlineExceeded",
            JsonException or InvalidOperationException => "UnityRequestFailed",
            _ => "UnityAttachFailed"
        };
        var text = error is OperationCanceledException ? "The Unity request deadline expired." : error.Message;
        text = Regex.Replace(text, @"(?i)(bearer\s+)[^\s,;]+", "$1[redacted]");
        text = Regex.Replace(text, "(?i)((?:api[_-]?key|token|password|secret)\\s*[=:]\\s*[\"']?)[^\\s,;\"']+", "$1[redacted]");
        text = DiagnosticPaths.Redact(text);
        return new JsonObject { ["code"] = code, ["message"] = text.Length > 4000 ? text[..4000] : text };
    }

    private static Request ParseRequest(string method, JsonObject? parameters)
    {
        if (parameters == null) throw new HostException("UnityParamsInvalid", "Request params must be an object.");
        var allowed = new HashSet<string> { "projectPath", "processId", "timeoutMs" };
        if (method == "execute") allowed.Add("code");
        if (parameters.Any(pair => !allowed.Contains(pair.Key))) throw new HostException("UnityParamsInvalid", "Unsupported request parameter.");
        var project = ReadString(parameters, "projectPath");
        int? pid = null;
        if (parameters.ContainsKey("processId"))
        {
            if (parameters["processId"] is not JsonValue pidValue || !pidValue.TryGetValue<int>(out var value) || value <= 0)
                throw new HostException("UnityParamsInvalid", "processId must be a positive integer.");
            pid = value;
        }
        var timeout = 60_000;
        if (parameters.ContainsKey("timeoutMs"))
        {
            if (parameters["timeoutMs"] is not JsonValue timeoutValue || !timeoutValue.TryGetValue<int>(out timeout) || timeout < 1 || timeout > 60_000)
                throw new HostException("UnityParamsInvalid", "timeoutMs must be between 1 and 60000.");
        }
        string? code = null;
        if (method == "execute")
        {
            code = ReadString(parameters, "code");
            if (string.IsNullOrWhiteSpace(code) || code.Contains('\0') || Encoding.UTF8.GetByteCount(code) > MaximumCodeBytes)
                throw new HostException("UnityCodeInvalid", "C# must be nonempty, contain no NUL, and fit in 128 KiB of UTF-8.");
        }
        return new Request(method, project, pid, timeout, code);
    }

    private static long ReadId(JsonObject request)
    {
        if (request["id"] is not JsonValue value || !value.TryGetValue<long>(out var id) || id < 1 || id > 9_007_199_254_740_991)
            throw new HostException("UnityIdInvalid", "Request id must be a positive JSON-safe integer.");
        return id;
    }
    private static string ReadString(JsonObject node, string name) => ReadOptionalString(node, name)
        ?? throw new HostException("UnityParamsInvalid", $"{name} must be a string.");
    private static string? ReadOptionalString(JsonObject node, string name) =>
        node[name] is JsonValue value && value.TryGetValue<string>(out var result) ? result : null;
    private sealed record Request(string Method, string ProjectPath, int? ProcessId, int TimeoutMs, string? Code);
    private sealed class DispatchState
    {
        private readonly object sync = new();
        private int sent;
        private bool closed;
        public bool Sent => Volatile.Read(ref sent) != 0;
        public void MarkSent(CancellationToken cancellationToken)
        {
            lock (sync)
            {
                if (closed) throw new OperationCanceledException(cancellationToken);
                cancellationToken.ThrowIfCancellationRequested();
                Interlocked.Exchange(ref sent, 1);
            }
        }
        public void Close() { lock (sync) closed = true; }
    }
    public async ValueTask DisposeAsync()
    {
        try { await backend.DisposeAsync().AsTask().WaitAsync(TimeSpan.FromSeconds(2)); }
        catch (Exception) { /* Host exit must not reconnect or block on an unavailable Editor. */ }
    }
}