wx-win-agent/node-agent/WxAgent.Core/RemoteTaskLedger.cs

using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using System.Text.Json.Serialization;

namespace WxAgent.Core;

public sealed record LocalRemoteTaskRecord
{
    [JsonPropertyName("taskId")]
    public string TaskId { get; init; } = "";

    [JsonPropertyName("accountId")]
    public string AccountId { get; init; } = "";

    [JsonPropertyName("leaseGeneration")]
    public long LeaseGeneration { get; init; }

    [JsonPropertyName("commandFingerprint")]
    public string CommandFingerprint { get; init; } = "";

    [JsonPropertyName("status")]
    public RemoteTaskStatus Status { get; init; }

    [JsonPropertyName("result")]
    public RemoteTaskResult? Result { get; init; }

    [JsonPropertyName("reportingScopes")]
    public IReadOnlyList<RemoteReportingScope> ReportingScopes { get; init; } = [];

    [JsonPropertyName("reported")]
    public bool Reported { get; init; }

    [JsonPropertyName("updatedAt")]
    public DateTimeOffset UpdatedAt { get; init; }
}

public sealed record UnreportedRemoteTask(
    RemoteTaskResult Result,
    IReadOnlyList<RemoteReportingScope> ReportingScopes);

public sealed class RemoteTaskLedger
{
    private readonly string _path;
    private readonly object _gate = new();
    private readonly Dictionary<string, LocalRemoteTaskRecord> _records;

    public RemoteTaskLedger(string path)
    {
        _path = Path.GetFullPath(path);
        _records = Load(_path);
        RecoverIncomplete();
    }

    public bool TryGet(string taskId, out LocalRemoteTaskRecord? record)
    {
        lock (_gate) return _records.TryGetValue(taskId, out record);
    }

    public bool Accept(RemoteTaskEnvelope task)
    {
        lock (_gate)
        {
            var fingerprint = Fingerprint(task);
            if (_records.TryGetValue(task.TaskId, out var existing))
            {
                if (!string.Equals(existing.CommandFingerprint, fingerprint, StringComparison.Ordinal)
                    || !string.Equals(existing.AccountId, task.AccountId, StringComparison.Ordinal))
                    throw new WxAgentException(WxAgentErrorCode.InvalidOperationState, "The same remote task ID was reused with different command parameters.");
                if (existing.Status is RemoteTaskStatus.Succeeded or RemoteTaskStatus.Failed or RemoteTaskStatus.Cancelled or RemoteTaskStatus.Expired or RemoteTaskStatus.ResultUnconfirmed)
                    return false;
                if (existing.LeaseGeneration != task.LeaseGeneration)
                {
                    _records[task.TaskId] = existing with
                    {
                        Status = RemoteTaskStatus.ResultUnconfirmed,
                        Result = new RemoteTaskResult(task.TaskId, task.AccountId, task.LeaseGeneration,
                            RemoteTaskStatus.ResultUnconfirmed, "LeaseGenerationChanged", "The node will not replay a task after its lease generation changed.", true, null, Guid.NewGuid().ToString("N")),
                        Reported = false,
                        UpdatedAt = DateTimeOffset.UtcNow
                    };
                    SaveLocked();
                    return false;
                }
                return true;
            }
            _records[task.TaskId] = new LocalRemoteTaskRecord
            {
                TaskId = task.TaskId,
                AccountId = task.AccountId,
                LeaseGeneration = task.LeaseGeneration,
                CommandFingerprint = fingerprint,
                Status = RemoteTaskStatus.Accepted,
                UpdatedAt = DateTimeOffset.UtcNow
            };
            SaveLocked();
            return true;
        }
    }

    public void MarkAccepted(RemoteTaskEnvelope task)
    {
        lock (_gate) UpdateLocked(task.TaskId, task.AccountId, task.LeaseGeneration, RemoteTaskStatus.Accepted, null, false);
    }

    public void MarkRunning(RemoteTaskEnvelope task)
    {
        lock (_gate) UpdateLocked(task.TaskId, task.AccountId, task.LeaseGeneration, RemoteTaskStatus.Running, null, false);
    }

    public void Complete(RemoteTaskResult result, IReadOnlyList<RemoteReportingScope>? chatScopes = null)
    {
        if (result.Status is not (RemoteTaskStatus.Succeeded or RemoteTaskStatus.Failed or RemoteTaskStatus.Cancelled or RemoteTaskStatus.Expired or RemoteTaskStatus.ResultUnconfirmed))
            throw new WxAgentException(WxAgentErrorCode.InvalidArgument, "Only a terminal result can be persisted in the remote task ledger.");
        lock (_gate)
        {
            if (!_records.TryGetValue(result.TaskId, out var existing))
            {
                existing = new LocalRemoteTaskRecord
                {
                    TaskId = result.TaskId,
                    AccountId = result.AccountId,
                    LeaseGeneration = result.LeaseGeneration,
                    CommandFingerprint = "unknown"
                };
            }
            if (!string.Equals(existing.AccountId, result.AccountId, StringComparison.Ordinal)
                || existing.LeaseGeneration > result.LeaseGeneration)
                throw new WxAgentException(WxAgentErrorCode.InvalidOperationState, "The remote task result does not match the local task ledger.");
            var scopes = chatScopes?.ToArray() ?? [];
            _records[result.TaskId] = existing with
            {
                AccountId = result.AccountId,
                LeaseGeneration = result.LeaseGeneration,
                Status = result.Status,
                Result = scopes.Length == 0 ? result with { Content = null } : result,
                ReportingScopes = scopes,
                Reported = false,
                UpdatedAt = DateTimeOffset.UtcNow
            };
            SaveLocked();
        }
    }

    public IReadOnlyList<RemoteTaskResult> UnreportedResults()
    {
        lock (_gate)
        {
            return _records.Values
                .Where(record => !record.Reported && record.Result is not null)
                .Select(record => record.Result!)
                .ToArray();
        }
    }

    public IReadOnlyList<UnreportedRemoteTask> UnreportedResultsWithScopes()
    {
        lock (_gate)
        {
            return _records.Values
                .Where(record => !record.Reported && record.Result is not null)
                .Select(record => new UnreportedRemoteTask(record.Result!, record.ReportingScopes))
                .ToArray();
        }
    }

    public void MarkReported(string taskId)
    {
        lock (_gate)
        {
            if (_records.TryGetValue(taskId, out var existing) && existing.Result is not null)
            {
                _records[taskId] = existing with { Reported = true, UpdatedAt = DateTimeOffset.UtcNow };
                SaveLocked();
            }
        }
    }

    private void RecoverIncomplete()
    {
        lock (_gate)
        {
            var changed = false;
            foreach (var pair in _records.ToArray())
            {
                if (pair.Value.Status is not (RemoteTaskStatus.Accepted or RemoteTaskStatus.Running)) continue;
                _records[pair.Key] = pair.Value with
                {
                    Status = RemoteTaskStatus.ResultUnconfirmed,
                    Result = new RemoteTaskResult(pair.Value.TaskId, pair.Value.AccountId, pair.Value.LeaseGeneration,
                        RemoteTaskStatus.ResultUnconfirmed, "AgentRestartedWithIncompleteTask", "Execution was not replayed.", true, null, Guid.NewGuid().ToString("N")),
                    Reported = false,
                    UpdatedAt = DateTimeOffset.UtcNow
                };
                changed = true;
            }
            if (changed) SaveLocked();
        }
    }

    private void UpdateLocked(string taskId, string accountId, long leaseGeneration, RemoteTaskStatus status, RemoteTaskResult? result, bool reported)
    {
        if (!_records.TryGetValue(taskId, out var existing))
            throw new WxAgentException(WxAgentErrorCode.InvalidOperationState, "The remote task was not accepted into the local ledger.");
        if (!string.Equals(existing.AccountId, accountId, StringComparison.Ordinal) || existing.LeaseGeneration != leaseGeneration)
            throw new WxAgentException(WxAgentErrorCode.InvalidOperationState, "The remote task lease does not match the local ledger.");
        _records[taskId] = existing with { Status = status, Result = result, Reported = reported, UpdatedAt = DateTimeOffset.UtcNow };
        SaveLocked();
    }

    private void SaveLocked()
    {
        var directory = Path.GetDirectoryName(_path) ?? AppContext.BaseDirectory;
        Directory.CreateDirectory(directory);
        var temporary = _path + ".tmp-" + Guid.NewGuid().ToString("N");
        try
        {
            File.WriteAllText(temporary, JsonSerializer.Serialize(_records.Values, RemoteJson.Options), Encoding.UTF8);
            if (!OperatingSystem.IsWindows())
            {
                try { File.SetUnixFileMode(temporary, UnixFileMode.UserRead | UnixFileMode.UserWrite); }
                catch (PlatformNotSupportedException) { }
            }
            if (OperatingSystem.IsWindows() && File.Exists(_path)) File.Replace(temporary, _path, null);
            else File.Move(temporary, _path, true);
        }
        finally
        {
            if (File.Exists(temporary)) File.Delete(temporary);
        }
    }

    private static Dictionary<string, LocalRemoteTaskRecord> Load(string path)
    {
        if (!File.Exists(path)) return new(StringComparer.Ordinal);
        try
        {
            var records = JsonSerializer.Deserialize<IReadOnlyList<LocalRemoteTaskRecord>>(File.ReadAllText(path), RemoteJson.Options) ?? [];
            return records.ToDictionary(record => record.TaskId, StringComparer.Ordinal);
        }
        catch (Exception exception) when (exception is IOException or JsonException or NotSupportedException or ArgumentException)
        {
            throw new WxAgentException(WxAgentErrorCode.InvalidOperationState, "Remote task ledger could not be loaded; remote task execution is blocked.", exception);
        }
    }

    private static string Fingerprint(RemoteTaskEnvelope task)
    {
        var value = $"{task.AccountId}\n{task.Kind}\n{task.Payload.GetRawText()}";
        return Convert.ToHexString(SHA256.HashData(Encoding.UTF8.GetBytes(value)));
    }
}