using Health.Application.AI; using Health.Infrastructure.Data.Records; namespace Health.Infrastructure.AI; public sealed class EfAiWriteConfirmationStore(AppDbContext db) : IAiWriteConfirmationStore { private readonly AppDbContext _db = db; public async Task CreateAsync( Guid userId, string toolName, string arguments, TimeSpan lifetime, CancellationToken ct) { var now = DateTime.UtcNow; if (IsInMemory()) { var expired = await _db.AiWriteCommands.Where(x => x.Status == "Pending" && x.ExpiresAt <= now).ToListAsync(ct); foreach (var item in expired) { item.Status = "Expired"; item.UpdatedAt = now; } } else { await _db.AiWriteCommands .Where(x => x.Status == "Pending" && x.ExpiresAt <= now) .ExecuteUpdateAsync(setters => setters .SetProperty(x => x.Status, "Expired") .SetProperty(x => x.UpdatedAt, now), ct); } var command = new AiWriteCommandRecord { Id = Guid.NewGuid(), UserId = userId, ToolName = toolName, Arguments = arguments, Status = "Pending", ExpiresAt = now.Add(lifetime), CreatedAt = now, UpdatedAt = now, }; await _db.AiWriteCommands.AddAsync(command, ct); await _db.SaveChangesAsync(ct); return ToCommand(command); } public async Task TakeAsync(Guid commandId, Guid userId, CancellationToken ct) { var now = DateTime.UtcNow; if (IsInMemory()) { var tracked = await _db.AiWriteCommands.FirstOrDefaultAsync( x => x.Id == commandId && x.UserId == userId && x.Status == "Pending" && x.ExpiresAt > now, ct); if (tracked == null) return null; tracked.Status = "Processing"; tracked.UpdatedAt = now; await _db.SaveChangesAsync(ct); } else { var updated = await _db.AiWriteCommands .Where(x => x.Id == commandId && x.UserId == userId && x.Status == "Pending" && x.ExpiresAt > now) .ExecuteUpdateAsync(setters => setters .SetProperty(x => x.Status, "Processing") .SetProperty(x => x.UpdatedAt, now), ct); if (updated == 0) return null; } var record = await _db.AiWriteCommands.AsNoTracking().FirstAsync(x => x.Id == commandId, ct); return ToCommand(record); } public async Task CompleteAsync(PendingAiWriteCommand command, CancellationToken ct) { if (IsInMemory()) { var tracked = await _db.AiWriteCommands.FirstOrDefaultAsync( x => x.Id == command.Id && x.UserId == command.UserId && x.Status == "Processing", ct); if (tracked != null) { tracked.Status = "Completed"; tracked.UpdatedAt = DateTime.UtcNow; await _db.SaveChangesAsync(ct); } return; } await _db.AiWriteCommands .Where(x => x.Id == command.Id && x.UserId == command.UserId && x.Status == "Processing") .ExecuteUpdateAsync(setters => setters .SetProperty(x => x.Status, "Completed") .SetProperty(x => x.UpdatedAt, DateTime.UtcNow), ct); } private bool IsInMemory() => _db.Database.ProviderName == "Microsoft.EntityFrameworkCore.InMemory"; private static PendingAiWriteCommand ToCommand(AiWriteCommandRecord record) => new( record.Id, record.UserId, record.ToolName, record.Arguments, record.ExpiresAt); }