using System.Reflection;
using System.Text.Json;
using AIStudio.Settings;
using AIStudio.Tools.AIJobs;
using AIStudio.Tools.DataBackup;
using AIStudio.Tools.Media;
using AIStudio.Tools.Metadata;
using AIStudio.Tools.PluginSystem;
using AIStudio.Tools.ToolCallingSystem;
namespace AIStudio.Tools.Services;
///
/// Prepares the restore of a backup, which the next start of AI Studio applies.
///
///
/// A restore replaces all data, so it runs only while nothing else changes it: no answer, no
/// transcript, and no backup may be running, and the indexing pauses from the start of the preparation
/// until the restart. Whatever changes after the preparation is lost with the restart, which is why
/// the app asks for the restart right away, and offers to discard the restore instead.
///
/// The service enforces what the organization allows on its own, since a component may show an
/// outdated state.
///
/// At the start which applied a restore, the service holds the report of the runtime until the user
/// saw it, cf. DataRestoreFollowUpService.
///
public sealed class DataRestoreService(ILogger logger, SettingsManager settingsManager, RustService rustService, ToolRegistry toolRegistry, ToolSettingsService toolSettingsService, DataSourceEmbeddingService embeddingService, DataBackupService dataBackupService, AIJobService aiJobService, MediaTranscriptionService mediaTranscriptionService)
{
private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(DataRestoreService).Namespace, nameof(DataRestoreService));
private readonly Lock stateLock = new();
private readonly TaskCompletionSource startupStateLoaded = new(TaskCreationOptions.RunContinuationsAsynchronously);
private CancellationTokenSource? preparing;
///
/// Holds the indexing back while a restore is prepared or waits for the restart.
///
private IAsyncDisposable? indexingPause;
///
/// Fires whenever the state changes: a preparation starts, makes progress, or ends, or a prepared restore is discarded.
///
///
/// It fires on whatever thread the preparation runs. A component hands it over to its own context.
///
public event Action? StateChanged;
///
/// Whether a restore is being prepared.
///
public bool IsPreparing
{
get
{
lock (this.stateLock)
return this.preparing is not null;
}
}
///
/// Whether a restore is prepared and waits for the restart.
///
public bool IsRestorePending { get; private set; }
///
/// How far the running preparation got.
///
public DataRestoreProgress Progress { get; private set; }
///
/// How the last preparation ended, or null when none ended yet.
///
public DataRestoreResult? LastResult { get; private set; }
///
/// How the runtime applied a restore at this start, until the user saw it; otherwise null.
///
public DataRestoreReport? Report { get; private set; }
///
/// Whether the runtime has not finished a restore and continues at the next start, cf.
/// DataRestoreFollowUp.IsRestoreUnfinished. Until then, AI Studio neither prepares nor discards a
/// restore.
///
public bool IsRestoreUnfinished { get; private set; }
///
/// Completes once the service knows what the last start left behind, cf. LoadStartupStateAsync.
///
public Task StartupStateLoaded => this.startupStateLoaded.Task;
///
/// Learns at the start what the runtime and the last run left behind: the report of a restore, a
/// restore which still waits for a restart, and what a backup or a preparation left halfway.
///
/// Cancels reading the report.
public async Task LoadStartupStateAsync(CancellationToken token)
{
try
{
var dataDirectory = SettingsManager.DataDirectory ?? throw new InvalidOperationException("The data directory is not known yet.");
var configDirectory = SettingsManager.ConfigDirectory ?? throw new InvalidOperationException("The configuration directory is not known yet.");
try
{
this.Report = await DataRestoreFollowUp.ReadReportAsync(dataDirectory, token);
if (this.Report is not null)
logger.LogInformation("The runtime applied a restore at this start: {Outcome}. {Detail}", this.Report.Outcome, this.Report.Detail);
}
catch (Exception exception) when (exception is JsonException or IOException)
{
logger.LogError(exception, "Could not read how the runtime applied a restore.");
DataRestoreFollowUp.RemoveReport(dataDirectory);
}
this.IsRestoreUnfinished = DataRestoreFollowUp.IsRestoreUnfinished(dataDirectory);
if (this.IsRestoreUnfinished)
logger.LogError("The runtime has not finished a restore. It tries to move the replaced data back at the next start.");
if (!dataBackupService.IsRunning && !this.IsPreparing)
DataRestoreFollowUp.RemoveLeftovers(dataDirectory, configDirectory);
//
// A restore which still waits for its start, e.g., in the development environment, where the
// .NET server may restart without the runtime: whatever changes until then is lost.
//
if (DataRestoreFollowUp.IsRestorePending(dataDirectory))
{
logger.LogWarning("A prepared restore still waits for the restart which applies it.");
this.IsRestorePending = true;
this.indexingPause ??= await embeddingService.PauseIndexingAsync(token);
}
}
catch (OperationCanceledException)
{
throw;
}
catch (Exception exception)
{
logger.LogError(exception, "Could not learn what the last start left behind of a restore or a backup.");
}
finally
{
this.startupStateLoaded.TrySetResult();
this.StateChanged?.Invoke();
}
}
///
/// Removes the data which a restore replaced. Call it once the restored settings loaded.
///
public void RemoveReplacedData()
{
try
{
var dataDirectory = SettingsManager.DataDirectory ?? throw new InvalidOperationException("The data directory is not known yet.");
var configDirectory = SettingsManager.ConfigDirectory ?? throw new InvalidOperationException("The configuration directory is not known yet.");
if (DataRestoreFollowUp.RemoveReplacedData(dataDirectory, configDirectory))
logger.LogInformation("Removed the data which the last restore replaced.");
}
catch (Exception exception)
{
logger.LogError(exception, "Could not remove the data which the last restore replaced.");
}
}
///
/// Finds what the user still has to do after the restore of this start, e.g., enter API keys again.
///
/// The open points; none when no restore was applied.
public async Task> CollectOpenPointsAsync()
{
if (this.Report?.Outcome is not DataRestoreOutcome.APPLIED)
return [];
try
{
return await DataRestoreOpenPoints.CollectAsync(settingsManager.ConfigurationData, new RustStoredSecrets(rustService), this.FindToolSecrets());
}
catch (Exception exception)
{
logger.LogError(exception, "Could not find what remains to do after the restore.");
return [];
}
}
///
/// Finds the secret settings of the tools which the user set up, i.e., saved the settings of.
///
///
/// A setting the organization fixed comes from its configuration plugin, and one the tool does
/// not show for the saved settings means nothing for them. Which of the remaining ones held a
/// value nobody knows: a backup holds no secrets, so the dialog asks to enter them again if they
/// were used.
///
private List FindToolSecrets()
{
var toolSecrets = new List();
foreach (var definition in toolRegistry.GetAllDefinitions())
{
if (toolRegistry.GetImplementation(definition.ImplementationKey) is not { } implementation)
continue;
if (!settingsManager.ConfigurationData.Tools.Settings.TryGetValue(definition.Id, out var storedValues))
continue;
foreach (var (fieldName, field) in definition.SettingsSchema.Properties)
if (field.Secret && !toolSettingsService.IsFieldLocked(definition, fieldName) && implementation.IsSettingsFieldVisible(fieldName, storedValues))
toolSecrets.Add(new(definition.Id, fieldName, implementation.GetDisplayName(), implementation.GetSettingsFieldLabel(fieldName, field)));
}
return toolSecrets;
}
///
/// Forgets the report of the runtime, once the user saw it.
///
public void ForgetReport()
{
this.Report = null;
try
{
DataRestoreFollowUp.RemoveReport(SettingsManager.DataDirectory ?? throw new InvalidOperationException("The data directory is not known yet."));
}
catch (Exception exception)
{
logger.LogError(exception, "Could not remove the report of the restore.");
}
this.StateChanged?.Invoke();
}
///
/// Prepares the restore of a backup for the next start.
///
/// The backup.
/// Whether the restore is prepared, or why not.
public async Task PrepareRestoreAsync(string archivePath)
{
if (!settingsManager.ConfigurationData.App.AllowUserToRestoreData)
return this.Refuse(TB("Your organization does not allow restoring backups."));
if (dataBackupService.IsRunning)
return this.Refuse(TB("A backup is running. Please restore once it is done."));
if (this.IsRestoreUnfinished)
return this.Refuse(TB("AI Studio has not finished restoring an earlier backup yet. Please restart AI Studio first."));
if (aiJobService.HasActiveJobs || mediaTranscriptionService.GetSnapshots().Any(snapshot => snapshot.Status is MediaImportStatus.QUEUED or MediaImportStatus.RUNNING or MediaImportStatus.CANCELING))
return this.Refuse(TB("AI Studio is still writing an answer or a transcript. Please restore once it is done."));
CancellationTokenSource cancellation;
lock (this.stateLock)
{
if (this.preparing is not null)
return new(false, false, TB("A restore is being prepared already."));
cancellation = this.preparing = new();
}
this.ReportProgress(new(0, 0));
var token = cancellation.Token;
try
{
var target = this.CreateTarget();
this.indexingPause ??= await embeddingService.PauseIndexingAsync(token);
var progress = new Progress(this.ReportProgress);
var result = await Task.Run(() => DataRestorePreparer.PrepareAsync(archivePath, target, progress, token), token);
this.IsRestorePending = result.Success;
if (result.Success)
logger.LogInformation("Prepared the restore of the backup '{ArchivePath}'. It is applied at the next start.", archivePath);
else
logger.LogWarning("Could not prepare the restore of the backup '{ArchivePath}': {Issue}", archivePath, result.Issue);
return this.End(result);
}
catch (OperationCanceledException) when (token.IsCancellationRequested)
{
logger.LogInformation("The user cancelled preparing the restore of the backup '{ArchivePath}'.", archivePath);
return this.End(new(false, true, string.Empty));
}
catch (Exception exception)
{
// The preparation removes what it unpacked when it fails:
logger.LogError(exception, "Preparing the restore of the backup '{ArchivePath}' failed.", archivePath);
return this.End(new(false, false, string.Format(TB("The restore could not be prepared: {0}"), exception.Message)));
}
finally
{
lock (this.stateLock)
this.preparing = null;
cancellation.Dispose();
if (!this.IsRestorePending)
await this.EndIndexingPauseAsync();
this.StateChanged?.Invoke();
}
}
///
/// Stops the running preparation. It removes what it unpacked so far.
///
public void Cancel()
{
lock (this.stateLock)
this.preparing?.Cancel();
}
///
/// Discards a prepared restore, so that the next start leaves all data as it is.
///
public async Task DiscardPendingRestoreAsync()
{
if (this.IsPreparing || this.IsRestoreUnfinished)
return;
try
{
var target = this.CreateTarget();
try
{
// The unpacked backup may hold thousands of files, and the click to discard it arrives on the thread of the user interface:
await Task.Run(() => DataRestorePreparer.Discard(target));
logger.LogInformation("Discarded the prepared restore.");
}
finally
{
// What remains on the disk decides, should discarding have failed halfway:
this.IsRestorePending = DataRestorePreparer.IsPrepared(target);
}
}
catch (Exception exception)
{
logger.LogError(exception, "Could not discard the prepared restore.");
}
// A restore which the next start still applies keeps the indexing paused:
if (!this.IsRestorePending)
await this.EndIndexingPauseAsync();
this.StateChanged?.Invoke();
}
private DataRestoreTarget CreateTarget()
{
var dataDirectory = SettingsManager.DataDirectory ?? throw new InvalidOperationException("The data directory is not known yet.");
var configDirectory = SettingsManager.ConfigDirectory ?? throw new InvalidOperationException("The configuration directory is not known yet.");
var appVersion = Assembly.GetExecutingAssembly().GetCustomAttribute()?.Version ?? string.Empty;
var organizationPluginFolders = DataBackupContents.FindOrganizationPluginsAmongUserPlugins(PluginFactory.AvailablePlugins.OfType(), dataDirectory);
// Loading the settings brings them to the version this AI Studio writes:
return new(dataDirectory, configDirectory, appVersion, settingsManager.ConfigurationData.Version, organizationPluginFolders);
}
private async Task EndIndexingPauseAsync()
{
var pause = this.indexingPause;
this.indexingPause = null;
if (pause is not null)
await pause.DisposeAsync();
}
private DataRestoreResult Refuse(string issue)
{
logger.LogWarning("Refused to prepare a restore: {Issue}", issue);
return this.End(new(false, false, issue));
}
private DataRestoreResult End(DataRestoreResult result)
{
this.LastResult = result;
this.StateChanged?.Invoke();
return result;
}
private void ReportProgress(DataRestoreProgress progress)
{
this.Progress = progress;
this.StateChanged?.Invoke();
}
}