MoShiro/backend/src/Monitor.Backend.Application/Tasks/TaskPlanner.cs

using System.Security.Cryptography;
using System.Text.Json;
using Fleet.Core;
using Fleet.Server;
using Monitor.Protocol;

namespace Monitor.Backend.Application;

public sealed record DesiredTaskSet(long Version, TaskSet Document);

/// <summary>
/// 为每个 Agent 计算期望状态（<see cref="TaskSet"/>）：有效任务集 + 允许连接的 Probe 公钥。
/// 只有内容真的变了才递增版本号，所以断线重连后 Agent 收到的永远是当前版本的完整快照。
/// </summary>
public sealed class TaskPlanner : IDisposable
{
    public const string ProbeAgentsVersionKey = "probe_agents_version";

    private readonly ITaskStore _tasks;
    private readonly IAgentStore _agents;
    private readonly IProbeStore _probes;
    private readonly ISettingsStore _settings;
    private readonly LinkService _links;
    private readonly SemaphoreSlim _gate = new(1, 1);

    public TaskPlanner(ITaskStore tasks, IAgentStore agents, IProbeStore probes, ISettingsStore settings, LinkService links)
    {
        _links = links;
        _tasks = tasks;
        _agents = agents;
        _probes = probes;
        _settings = settings;
    }

    /// <summary>pending、disabled 或不存在的 Agent 返回 null（不下发状态）。</summary>
    public async ValueTask<DesiredTaskSet?> GetDesiredAsync(NodeId agentId, CancellationToken ct)
    {
        AgentRecord? agent = await _agents.GetAsync(agentId, ct);
        if (agent is null || agent.Status != NodeStatus.Active)
        {
            return null;
        }

        IReadOnlyList<TaskRecord> all = await _tasks.ListAsync(ct);
        IReadOnlyList<ProbeRecord> probes = await _probes.ListAsync(ct);
        LinkSnapshot links = await _links.SnapshotAsync(ct);

        var document = new TaskSet();
        foreach (EffectiveTask task in EffectiveTasks.Compute(all, agent, filterByAgentKinds: true))
        {
            document.Tasks.Add(new TaskSpec
            {
                Id = task.Id,
                Scope = task.Scope,
                Kind = task.Kind,
                Target = task.Target,
                Params = new Dictionary<string, JsonElement>(task.Params),
                Schedule = task.Schedule,
                Enabled = true,
            });
        }

        foreach (ProbeRecord probe in probes.Where(p => p.RevokedAtMs is null && links.IsLinked(p.Id, agentId.Value)))
        {
            document.Peers.Add(new PeerSpec { ProbeId = new ProbeId(probe.Id), PublicKey = Convert.ToBase64String(probe.PublicKey) });
        }

        string hash = Convert.ToHexStringLower(SHA256.HashData(JsonSerializer.SerializeToUtf8Bytes(document, MonitorJsonContext.Default.TaskSet)));

        await _gate.WaitAsync(ct);
        try
        {
            // 重新读取，避免并发的两次计算各自递增
            AgentRecord current = await _agents.GetAsync(agentId, ct) ?? agent;
            string? storedHash = await _agents.GetTaskHashAsync(agentId, ct);
            long version = current.TaskVersion;
            if (storedHash != hash)
            {
                version++;
                await _agents.SetTaskVersionAsync(agentId, version, hash, ct);
            }

            document.Version = version;
            return new DesiredTaskSet(version, document);
        }
        finally
        {
            _gate.Release();
        }
    }

    public void Dispose() => _gate.Dispose();

    public ValueTask<long> BumpProbeAgentsVersionAsync(CancellationToken ct) => _settings.IncrementAsync(ProbeAgentsVersionKey, ct);
}