zcbot/windows-node/Zcbot.WindowsNode/OriginWorkerRunner.cs

using System.Collections.Concurrent;
using System.Diagnostics;
using System.Text;
using System.Text.Json;

namespace Zcbot.WindowsNode;

internal sealed class OriginWorkerRunner(JobInboxStore inbox)
{
    private static readonly TimeSpan WorkerTimeout = TimeSpan.FromMinutes(30);
    private readonly ConcurrentDictionary<Guid, Task> active = new();
    private readonly ConcurrentDictionary<Guid, CancellationTokenSource> cancellations = new();

    internal Task RunAsync(RecoverableJob job) =>
        active.GetOrAdd(job.JobId, _ => RunOnceAsync(job, CancellationFor(job.JobId).Token));

    internal void Cancel(Guid jobId)
    {
        CancellationFor(jobId).Cancel();
    }

    private CancellationTokenSource CancellationFor(Guid jobId) =>
        cancellations.GetOrAdd(jobId, _ => new CancellationTokenSource());

    private async Task RunOnceAsync(RecoverableJob job, CancellationToken cancellationToken)
    {
        try
        {
            var paths = NodePaths.ForCurrentMachine();
            var jobDirectory = Path.Combine(paths.JobsDirectory, job.JobId.ToString("D"));
            var terminalPath = Path.Combine(jobDirectory, "terminal.json");
            if (File.Exists(terminalPath)) return;

            var markerPath = Path.Combine(jobDirectory, "worker-started.json");
            if (File.Exists(markerPath))
            {
                inbox.WriteTerminal(
                    job,
                    "failed",
                    "NODE_RESTARTED_DURING_JOB",
                    "The node restarted after Origin execution began and cannot prove the prior worker state.");
                return;
            }

            var interpreter = OriginWorkerRuntime.ResolveInterpreter()
                ?? throw new InvalidOperationException("The fixed Origin Python interpreter is unavailable.");
            var workerScript = Path.GetFullPath(
                Path.Combine(AppContext.BaseDirectory, "origin-worker", "worker.py"));
            if (!File.Exists(workerScript))
            {
                throw new FileNotFoundException("The fixed Origin worker script is missing.", workerScript);
            }
            WriteMarker(markerPath, interpreter, workerScript);

            var startInfo = new ProcessStartInfo
            {
                FileName = interpreter,
                WorkingDirectory = jobDirectory,
                UseShellExecute = false,
                CreateNoWindow = true,
                RedirectStandardOutput = true,
                RedirectStandardError = true,
                StandardOutputEncoding = Encoding.UTF8,
                StandardErrorEncoding = Encoding.UTF8,
            };
            startInfo.ArgumentList.Add(workerScript);
            startInfo.ArgumentList.Add(jobDirectory);
            using var process = Process.Start(startInfo)
                ?? throw new InvalidOperationException("The fixed Origin worker did not start.");
            var stdout = process.StandardOutput.ReadToEndAsync();
            var stderr = process.StandardError.ReadToEndAsync();
            using var timeout = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
            timeout.CancelAfter(WorkerTimeout);
            try
            {
                await process.WaitForExitAsync(timeout.Token);
            }
            catch (OperationCanceledException)
            {
                process.Kill(entireProcessTree: true);
                if (cancellationToken.IsCancellationRequested)
                {
                    inbox.WriteTerminal(job, "cancelled", "USER_CANCELLED", "Cancelled by user.");
                }
                else
                {
                    inbox.WriteTerminal(job, "failed", "ORIGIN_WORKER_TIMEOUT", "Origin worker exceeded 30 minutes.");
                }
                return;
            }
            var output = await stdout;
            var error = await stderr;
            WriteDiagnostic(jobDirectory, output, error, process.ExitCode);
            if (!File.Exists(terminalPath))
            {
                inbox.WriteTerminal(
                    job,
                    "failed",
                    "ORIGIN_WORKER_NO_TERMINAL",
                    $"Origin worker exited with code {process.ExitCode} without terminal.json.");
            }
        }
        catch (Exception exception) when (
            exception is IOException
                or JsonException
                or UnauthorizedAccessException
                or InvalidOperationException)
        {
            inbox.WriteTerminal(job, "failed", "ORIGIN_WORKER_START_FAILED", exception.Message[..Math.Min(500, exception.Message.Length)]);
        }
        finally
        {
            active.TryRemove(job.JobId, out _);
            if (cancellations.TryRemove(job.JobId, out var cancellation)) cancellation.Dispose();
        }
    }

    private static void WriteMarker(string path, string interpreter, string workerScript)
    {
        var value = JsonSerializer.SerializeToUtf8Bytes(new
        {
            started_at = DateTimeOffset.UtcNow,
            node_pid = Environment.ProcessId,
            interpreter,
            worker_script = workerScript,
        });
        using var stream = new FileStream(
            path, FileMode.CreateNew, FileAccess.Write, FileShare.None,
            bufferSize: 4096, FileOptions.WriteThrough);
        stream.Write(value);
        stream.Flush(flushToDisk: true);
    }

    private static void WriteDiagnostic(string jobDirectory, string output, string error, int exitCode)
    {
        var logs = Path.Combine(jobDirectory, "logs");
        Directory.CreateDirectory(logs);
        var value = JsonSerializer.Serialize(new
        {
            exit_code = exitCode,
            stdout = output[..Math.Min(output.Length, 16 * 1024)],
            stderr = error[..Math.Min(error.Length, 16 * 1024)],
        });
        File.WriteAllText(Path.Combine(logs, "worker-process.json"), value, Encoding.UTF8);
    }
}

internal static class OriginWorkerRuntime
{
    internal static string? ResolveInterpreter()
    {
        var paths = NodePaths.ForCurrentMachine();
        var configured = Environment.GetEnvironmentVariable("ZCBOT_ORIGIN_PYTHON");
        var candidate = string.IsNullOrWhiteSpace(configured)
            ? Path.Combine(paths.RootDirectory, "runtimes", "origin", "python.exe")
            : configured;
        if (!Path.IsPathFullyQualified(candidate)) return null;
        var resolved = Path.GetFullPath(candidate);
        return File.Exists(resolved)
            && Path.GetFileName(resolved).Equals("python.exe", StringComparison.OrdinalIgnoreCase)
            ? resolved
            : null;
    }
}