YMhut-box-C-/src/box-winUI/Services/ProcessToolWorkerService.cs

using System.Diagnostics;
using System.IO.Pipes;
using System.Text;
using YMhut.Box.Core.Api;
using YMhut.Box.Core.Data;
using YMhut.Box.Core.Logging;
using YMhut.Box.Core.Tools;

namespace YMhut.Box.WinUI.Services;

public sealed class ProcessToolWorkerService : IToolWorkerService, IDisposable
{
    private readonly IApiManager? _apiManager;
    private readonly IReferenceDataService? _referenceDataService;
    private readonly ILogService? _logService;
    private readonly SemaphoreSlim _requestGate = new(1, 1);
    private readonly SemaphoreSlim _writeGate = new(1, 1);

    private Process? _process;
    private NamedPipeServerStream? _pipe;
    private StreamReader? _reader;
    private StreamWriter? _writer;
    private bool _disposed;

    public ProcessToolWorkerService(
        IApiManager? apiManager = null,
        IReferenceDataService? referenceDataService = null,
        ILogService? logService = null)
    {
        _apiManager = apiManager;
        _referenceDataService = referenceDataService;
        _logService = logService;
        AppDomain.CurrentDomain.ProcessExit += (_, _) => Dispose();
    }

    public Task<T> RunAsync<T>(Func<CancellationToken, Task<T>> work, CancellationToken cancellationToken = default)
    {
        return Task.Run(() => work(cancellationToken), cancellationToken);
    }

    public async Task<ToolExecutionResult> ExecuteToolAsync(
        IToolModule module,
        string input,
        IApiManager? apiManager = null,
        IReferenceDataService? referenceDataService = null,
        CancellationToken cancellationToken = default,
        string language = "zh-CN")
    {
        await _requestGate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            return await ExecuteThroughWorkerAsync(module, input, language, cancellationToken).ConfigureAwait(false);
        }
        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
        {
            throw;
        }
        catch (Exception exception)
        {
            await WriteWorkerLogAsync("Warning", "worker", "Tool worker unavailable; using in-process fallback", exception.Message, CancellationToken.None)
                .ConfigureAwait(false);
            return await ToolExecutor.ExecuteAsync(
                module,
                input,
                cancellationToken,
                apiManager ?? _apiManager,
                referenceDataService ?? _referenceDataService,
                language).ConfigureAwait(false);
        }
        finally
        {
            _requestGate.Release();
        }
    }

    public void Dispose()
    {
        if (_disposed)
        {
            return;
        }

        _disposed = true;
        try
        {
            if (_writer is not null)
            {
                _writer.WriteLine(ToolWorkerProtocol.Serialize(new ToolWorkerMessage(ToolWorkerProtocol.Shutdown)));
                _writer.Flush();
            }
        }
        catch
        {
        }

        DisposeWorker(killProcess: false);
        _requestGate.Dispose();
        _writeGate.Dispose();
    }

    private async Task<ToolExecutionResult> ExecuteThroughWorkerAsync(
        IToolModule module,
        string input,
        string language,
        CancellationToken cancellationToken)
    {
        await EnsureWorkerAsync(cancellationToken).ConfigureAwait(false);

        var requestId = Guid.NewGuid().ToString("N");
        var request = new ToolWorkerMessage(
            ToolWorkerProtocol.ExecuteTool,
            requestId,
            ToolId: module.Id,
            Input: input,
            TimeoutMs: 120_000,
            Language: language);

        await WriteMessageAsync(request, cancellationToken).ConfigureAwait(false);
        await WriteWorkerLogAsync("Information", "worker", $"Queued tool in worker: {module.Id}", null, cancellationToken)
            .ConfigureAwait(false);

        using var cancellation = cancellationToken.Register(() =>
        {
            _ = Task.Run(async () =>
            {
                try
                {
                    await WriteMessageAsync(
                        new ToolWorkerMessage(ToolWorkerProtocol.Cancel, requestId),
                        CancellationToken.None).ConfigureAwait(false);
                }
                catch
                {
                }
            });
        });

        while (true)
        {
            var message = await ReadMessageAsync(cancellationToken).ConfigureAwait(false);
            if (!string.Equals(message.RequestId, requestId, StringComparison.Ordinal))
            {
                continue;
            }

            if (string.Equals(message.Type, ToolWorkerProtocol.Result, StringComparison.Ordinal))
            {
                return new ToolExecutionResult(message.Ok, message.Output ?? string.Empty, message.Error, message.Document);
            }

            if (string.Equals(message.Type, ToolWorkerProtocol.Error, StringComparison.Ordinal))
            {
                return ToolExecutionResult.Fail(message.Error ?? "Worker failed to execute the tool.");
            }
        }
    }

    private async Task EnsureWorkerAsync(CancellationToken cancellationToken)
    {
        if (_process is { HasExited: false } && _pipe?.IsConnected == true && _reader is not null && _writer is not null)
        {
            return;
        }

        DisposeWorker(killProcess: true);

        var executable = ResolveWorkerExecutable()
            ?? throw new FileNotFoundException("YMhut.Box.Worker.exe was not found in the application output.");
        var pipeName = $"YMhutBoxWorker-{Environment.ProcessId}-{Guid.NewGuid():N}";
        _pipe = new NamedPipeServerStream(
            pipeName,
            PipeDirection.InOut,
            1,
            PipeTransmissionMode.Byte,
            PipeOptions.Asynchronous);

        var startInfo = new ProcessStartInfo
        {
            FileName = executable,
            Arguments = $"--pipe {pipeName}",
            WorkingDirectory = Path.GetDirectoryName(executable) ?? AppContext.BaseDirectory,
            UseShellExecute = false,
            CreateNoWindow = true
        };

        _process = Process.Start(startInfo)
            ?? throw new InvalidOperationException("Unable to start the YMhut tool worker process.");

        await _pipe.WaitForConnectionAsync(cancellationToken).WaitAsync(TimeSpan.FromSeconds(8), cancellationToken)
            .ConfigureAwait(false);
        _reader = new StreamReader(_pipe, Encoding.UTF8, detectEncodingFromByteOrderMarks: false, bufferSize: 4096, leaveOpen: true);
        _writer = new StreamWriter(_pipe, new UTF8Encoding(encoderShouldEmitUTF8Identifier: false), bufferSize: 4096, leaveOpen: true)
        {
            AutoFlush = true
        };

        var ready = await ReadMessageAsync(cancellationToken).ConfigureAwait(false);
        if (!string.Equals(ready.Type, ToolWorkerProtocol.Ready, StringComparison.Ordinal) ||
            !string.Equals(ready.Version, ToolWorkerProtocol.Version, StringComparison.Ordinal))
        {
            throw new InvalidOperationException("Tool worker did not complete the protocol handshake.");
        }

        await WriteWorkerLogAsync("Information", "worker", "Tool worker process started", Path.GetFileName(executable), cancellationToken)
            .ConfigureAwait(false);
    }

    private async Task WriteMessageAsync(ToolWorkerMessage message, CancellationToken cancellationToken)
    {
        var writer = _writer ?? throw new InvalidOperationException("Tool worker pipe is not connected.");
        await _writeGate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            await writer.WriteLineAsync(ToolWorkerProtocol.Serialize(message).AsMemory(), cancellationToken).ConfigureAwait(false);
            await writer.FlushAsync(cancellationToken).ConfigureAwait(false);
        }
        finally
        {
            _writeGate.Release();
        }
    }

    private async Task<ToolWorkerMessage> ReadMessageAsync(CancellationToken cancellationToken)
    {
        var reader = _reader ?? throw new InvalidOperationException("Tool worker pipe is not connected.");
        var line = await reader.ReadLineAsync(cancellationToken).ConfigureAwait(false);
        if (line is null)
        {
            DisposeWorker(killProcess: true);
            throw new IOException("Tool worker pipe closed unexpectedly.");
        }

        return ToolWorkerProtocol.Deserialize(line)
            ?? throw new InvalidDataException("Tool worker returned an invalid message.");
    }

    private void DisposeWorker(bool killProcess)
    {
        try
        {
            _reader?.Dispose();
            _writer?.Dispose();
            _pipe?.Dispose();
        }
        catch
        {
        }
        finally
        {
            _reader = null;
            _writer = null;
            _pipe = null;
        }

        try
        {
            if (_process is not null && !_process.HasExited && killProcess)
            {
                _process.Kill(entireProcessTree: true);
            }
        }
        catch
        {
        }
        finally
        {
            _process?.Dispose();
            _process = null;
        }
    }

    private async Task WriteWorkerLogAsync(
        string level,
        string category,
        string message,
        string? detail,
        CancellationToken cancellationToken)
    {
        if (_logService is null)
        {
            return;
        }

        await _logService.WriteAsync(level, category, message, detail, cancellationToken).ConfigureAwait(false);
    }

    private static string? ResolveWorkerExecutable()
    {
        var baseDirectory = AppContext.BaseDirectory;
        var candidates = new List<string>
        {
            Path.Combine(baseDirectory, "worker", "win-x64", "YMhut.Box.Worker.exe"),
            Path.Combine(baseDirectory, "worker", "YMhut.Box.Worker.exe"),
            Path.Combine(baseDirectory, "YMhut.Box.Worker.exe")
        };

        var directory = new DirectoryInfo(baseDirectory);
        while (directory is not null)
        {
            candidates.Add(Path.Combine(directory.FullName, "src", "YMhut.Box.Worker", "bin", "Debug", "net10.0", "YMhut.Box.Worker.exe"));
            candidates.Add(Path.Combine(directory.FullName, "src", "YMhut.Box.Worker", "bin", "Release", "net10.0", "YMhut.Box.Worker.exe"));
            directory = directory.Parent;
        }

        return candidates.FirstOrDefault(File.Exists);
    }
}