AI-Studio/app/MindWork AI Studio/Tools/Services/ProvenanceCleanupService.cs

using AIStudio.Assistants.BeyondFactsLab;
using AIStudio.Chat;
using AIStudio.Settings;
using AIStudio.Tools.DataBackup;
using AIStudio.Tools.Databases;
using AIStudio.Tools.Databases.ProvenanceStore;

namespace AIStudio.Tools.Services;

/// <summary>
/// One-shot startup service that removes from the provenance store what no stored chat refers to, see ProvenanceCleanup.
/// </summary>
/// <remarks>
/// It knows the chats of the workspaces, the temporary chats, and the chats of the BeyondFacts Lab. It reads the
/// latter also while the lab is switched off: otherwise, they would lose what they received at every start.
/// </remarks>
public sealed class ProvenanceCleanupService(DatabaseClientProvider databaseClientProvider, DataRestoreService dataRestoreService, BeyondFactsLabStore labStore, ILogger<ProvenanceCleanupService> logger) : BackgroundService
{
    /// <summary>
    /// When this run of the app started. Entries stored since belong to it and stay.
    /// </summary>
    private readonly DateTimeOffset startedAt = DateTimeOffset.UtcNow;

    /// <summary>Waits for the data directory, performs one cleanup pass, and then exits.</summary>
    /// <param name="stoppingToken">Host shutdown token.</param>
    protected override async Task ExecuteAsync(CancellationToken stoppingToken)
    {
        while (string.IsNullOrWhiteSpace(SettingsManager.DataDirectory) && !stoppingToken.IsCancellationRequested)
            await Task.Delay(TimeSpan.FromSeconds(1), stoppingToken);

        if (stoppingToken.IsCancellationRequested)
            return;

        //
        // A restore which the runtime could not finish leaves chats and store of two different states
        // in the data folder, while the chats it replaced wait elsewhere to be moved back. Nothing
        // tells then which chat refers to what:
        //
        if (SettingsManager.DataDirectory is { } dataDirectory && DataRestoreFollowUp.IsRestoreUnfinished(dataDirectory))
        {
            logger.LogWarning("A restore is not finished yet, so the provenance startup cleanup waits for a later start.");
            return;
        }

        try
        {
            var store = await databaseClientProvider.GetProvenanceStoreAsync(stoppingToken);
            if (store is NoProvenanceStoreClient)
            {
                logger.LogInformation("The provenance store is not available, so its startup cleanup has nothing to do.");
                return;
            }

            //
            // A restore at this start may have brought back the chats of one moment and kept the store
            // of another, e.g., from a backup without the store. Then the time of the last cleanup
            // tells nothing about the chats, and every chat is read:
            //
            await dataRestoreService.StartupStateLoaded.WaitAsync(stoppingToken);
            var lastCleanup = dataRestoreService.Report is null ? await store.GetLastCleanupAsync(stoppingToken) : null;

            //
            // SQLite works synchronously behind its async methods, and every changed chat with entries
            // is read once, so the whole pass runs on a thread of its own:
            //
            var result = await Task.Run(() => ProvenanceCleanup.CleanUpAsync(
                store,
                this.ReadAllChatLocations(),
                this.LoadChatAsync,
                this.GetChatWrittenAt,
                this.startedAt,
                lastCleanup,
                stoppingToken), stoppingToken);

            logger.LogInformation("The provenance startup cleanup removed {OrphanedCount} entries of chats which no longer exist and {UnreferencedCount} entries their chats no longer refer to. It skipped {SkippedCount} unchanged chats.", result.Orphaned, result.Unreferenced, result.Skipped);

            // A chat which could not be read is read again at the next start, as long as the time of this cleanup is not noted:
            if (result.Unreadable is 0)
                await store.SetLastCleanupAsync(this.startedAt, stoppingToken);
            else
                logger.LogWarning("The provenance startup cleanup could not read {UnreadableCount} chats, which keep all their entries. The next start reads them again.", result.Unreadable);
        }
        catch (OperationCanceledException)
        {
            // The app is shutting down; the next start cleans up.
        }
        catch (Exception exception)
        {
            logger.LogWarning(exception, "The provenance startup cleanup failed. The next start tries again.");
        }
    }

    /// <summary>
    /// Reads where every chat on the disk lies, from the folders alone.
    /// </summary>
    /// <returns>The place of every chat, by its ID.</returns>
    private IReadOnlyDictionary<Guid, StoredChatLocation> ReadAllChatLocations()
    {
        var locations = new Dictionary<Guid, StoredChatLocation>();
        foreach (var (chatId, workspaceId) in WorkspaceBehaviour.ReadAllChatLocations())
            locations[chatId] = StoredChatLocation.InWorkspace(workspaceId);

        foreach (var (chatId, labThread) in labStore.ReadAllThreadLocations())
            locations[chatId] = StoredChatLocation.InLab(labThread);

        return locations;
    }

    /// <summary>
    /// Reads a chat from its place.
    /// </summary>
    /// <param name="location">The place of the chat.</param>
    /// <param name="chatId">The ID of the chat.</param>
    /// <returns>The chat, or null when it cannot be read.</returns>
    private async Task<ChatThread?> LoadChatAsync(StoredChatLocation location, Guid chatId)
    {
        if (location.LabThread is not { } labThread)
            return await WorkspaceBehaviour.LoadChatAsync(new(location.WorkspaceId, chatId));

        try
        {
            return await labStore.LoadThreadAsync(labThread);
        }
        catch (Exception exception) when (exception is IOException or UnauthorizedAccessException or InvalidDataException)
        {
            // A chat which cannot be read keeps everything it received:
            logger.LogDebug(exception, "The provenance startup cleanup could not read a chat of the BeyondFacts Lab.");
            return null;
        }
    }

    /// <summary>
    /// Tells when the file of a chat was last written.
    /// </summary>
    /// <param name="location">The place of the chat.</param>
    /// <param name="chatId">The ID of the chat.</param>
    /// <returns>The time, or null when it is not known.</returns>
    private DateTimeOffset? GetChatWrittenAt(StoredChatLocation location, Guid chatId) => location.LabThread is { } labThread
        ? labStore.GetThreadWrittenAt(labThread)
        : WorkspaceBehaviour.GetChatWrittenAt(location.WorkspaceId, chatId);
}