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(); } }