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 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 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 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 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 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 { "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(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(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(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(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. */ } } }