Control-Center/backend/Repositories/AgentStateRepository.cs

using ControlCenter.Api.Data;
using ControlCenter.Api.Entities;
using ControlCenter.Api.Models;
using Microsoft.EntityFrameworkCore;

namespace ControlCenter.Api.Repositories;

/// <summary>
/// EF Core implementation of the Agent State Repository.
/// Maps between the persisted <see cref="Agent"/> entity and the
/// read-oriented <see cref="AgentState"/> model.
/// </summary>
public class AgentStateRepository : IAgentStateRepository
{
    private readonly AppDbContext _db;

    public AgentStateRepository(AppDbContext db)
    {
        _db = db;
    }

    /// <inheritdoc />
    public async Task<IReadOnlyList<AgentState>> GetAllAsync(CancellationToken ct = default)
    {
        var agents = await _db.Agents
            .AsNoTracking()
            .OrderByDescending(a => a.LastActivity)
            .ToListAsync(ct);

        return agents.Select(ToModel).ToList();
    }

    /// <inheritdoc />
    public async Task<AgentState?> GetBySessionKeyAsync(string sessionKey, CancellationToken ct = default)
    {
        var agent = await _db.Agents
            .AsNoTracking()
            .FirstOrDefaultAsync(a => a.SessionKey == sessionKey, ct);

        return agent is null ? null : ToModel(agent);
    }

    /// <inheritdoc />
    public async Task<bool> UpdateStatusAsync(Guid id, string status, CancellationToken ct = default)
    {
        if (!Enum.TryParse<AgentStatus>(status, ignoreCase: true, out var parsedStatus))
            return false;

        var agent = await _db.Agents.FindAsync([id], ct);
        if (agent is null)
            return false;

        agent.Status = parsedStatus;
        agent.UpdatedAt = DateTime.UtcNow;
        agent.LastActivity = DateTime.UtcNow;

        await _db.SaveChangesAsync(ct);
        return true;
    }

    /// <summary>
    /// Maps a persisted <see cref="Agent"/> entity to a <see cref="AgentState"/> model.
    /// </summary>
    private static AgentState ToModel(Agent agent) => new()
    {
        Id = agent.Id,
        Status = agent.Status.ToString(),
        Task = agent.Task,
        Progress = agent.Progress,
        SessionKey = agent.SessionKey,
        Channel = agent.Channel,
        LastActivity = agent.LastActivity,
        CreatedAt = agent.CreatedAt,
        UpdatedAt = agent.UpdatedAt,
    };
}