diff --git a/StabilityMatrix.Avalonia/ViewModels/Base/PausableProgressItemViewModelBase.cs b/StabilityMatrix.Avalonia/ViewModels/Base/PausableProgressItemViewModelBase.cs index 0881e65d..695da5b6 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Base/PausableProgressItemViewModelBase.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Base/PausableProgressItemViewModelBase.cs @@ -14,7 +14,9 @@ public abstract partial class PausableProgressItemViewModelBase : ProgressItemVi nameof(IsPaused), nameof(IsCompleted), nameof(CanPauseResume), - nameof(CanCancel) + nameof(CanCancel), + nameof(CanRetry), + nameof(CanDismiss) )] private ProgressState state = ProgressState.Inactive; @@ -33,9 +35,31 @@ public abstract partial class PausableProgressItemViewModelBase : ProgressItemVi public virtual bool SupportsPauseResume => true; public virtual bool SupportsCancel => true; + /// + /// Override to true in subclasses that support manual retry after failure. + /// Defaults to false so unrelated progress item types are never affected. + /// + public virtual bool SupportsRetry => false; + + /// + /// Override to true in subclasses that support dismissing a failed item, + /// which runs full sidecar cleanup before removing the entry. + /// + public virtual bool SupportsDismiss => false; + public bool CanPauseResume => SupportsPauseResume && !IsCompleted && !IsPending; public bool CanCancel => SupportsCancel && !IsCompleted; + /// + /// True only when this item supports retry AND is in the Failed state. + /// + public bool CanRetry => SupportsRetry && State == ProgressState.Failed; + + /// + /// True only when this item supports dismiss AND is in the Failed state. + /// + public bool CanDismiss => SupportsDismiss && State == ProgressState.Failed; + private AsyncRelayCommand? pauseCommand; public IAsyncRelayCommand PauseCommand => pauseCommand ??= new AsyncRelayCommand(Pause); @@ -51,6 +75,16 @@ public abstract partial class PausableProgressItemViewModelBase : ProgressItemVi public virtual Task Cancel() => Task.CompletedTask; + private AsyncRelayCommand? retryCommand; + public IAsyncRelayCommand RetryCommand => retryCommand ??= new AsyncRelayCommand(Retry); + + public virtual Task Retry() => Task.CompletedTask; + + private AsyncRelayCommand? dismissCommand; + public IAsyncRelayCommand DismissCommand => dismissCommand ??= new AsyncRelayCommand(Dismiss); + + public virtual Task Dismiss() => Task.CompletedTask; + [RelayCommand] private Task TogglePauseResume() { diff --git a/StabilityMatrix.Avalonia/ViewModels/Progress/DownloadProgressItemViewModel.cs b/StabilityMatrix.Avalonia/ViewModels/Progress/DownloadProgressItemViewModel.cs index 04809ec2..c763d5a4 100644 --- a/StabilityMatrix.Avalonia/ViewModels/Progress/DownloadProgressItemViewModel.cs +++ b/StabilityMatrix.Avalonia/ViewModels/Progress/DownloadProgressItemViewModel.cs @@ -71,6 +71,17 @@ private void OnProgressStateChanged(ProgressState state) } } + /// + /// Downloads support manual retry when they reach the Failed state. + /// + public override bool SupportsRetry => true; + + /// + /// Downloads support dismiss, which cleans up all sidecar files when + /// the user discards a failed download without retrying. + /// + public override bool SupportsDismiss => true; + /// public override Task Cancel() { @@ -91,4 +102,23 @@ public override Task Resume() { return downloadService.TryResumeDownload(download); } + + /// + /// Resets the internal retry counter so the user gets a fresh 3-attempt budget, + /// then re-registers the download in the service dictionary (it was removed on + /// failure) and resumes it through the normal concurrency queue. + public override Task Retry() + { + download.ResetAttempts(); + return downloadService.TryRestartDownload(download); + } + + /// + /// Runs full cleanup (temp file + sidecar files) for a failed download the user + /// chooses not to retry, then transitions to Cancelled so the service removes it. + public override Task Dismiss() + { + download.Dismiss(); + return Task.CompletedTask; + } } diff --git a/StabilityMatrix.Avalonia/Views/ProgressManagerPage.axaml b/StabilityMatrix.Avalonia/Views/ProgressManagerPage.axaml index c2286200..87dfbadc 100644 --- a/StabilityMatrix.Avalonia/Views/ProgressManagerPage.axaml +++ b/StabilityMatrix.Avalonia/Views/ProgressManagerPage.axaml @@ -113,6 +113,24 @@ IsVisible="{Binding CanCancel}"> + + + + + + ? ProgressUpdate; @@ -119,6 +122,13 @@ private void EnsureDownloadService() } } + private void CancelRetryDelay() + { + retryDelayCancellationTokenSource?.Cancel(); + retryDelayCancellationTokenSource?.Dispose(); + retryDelayCancellationTokenSource = null; + } + private async Task StartDownloadTask(long resumeFromByte, CancellationToken cancellationToken) { var progress = new Progress(OnProgressUpdate); @@ -184,6 +194,9 @@ internal void Start() $"Download state must be inactive or pending to start, not {ProgressState}" ); } + // Cancel any pending auto-retry delay (defensive: Start() accepts Inactive state). + CancelRetryDelay(); + Logger.Debug("Starting download {Download}", FileName); EnsureDownloadService(); @@ -201,13 +214,21 @@ internal void Start() internal void Resume() { - if (ProgressState != ProgressState.Inactive && ProgressState != ProgressState.Paused) + // Cancel any pending auto-retry delay since we're resuming now. + CancelRetryDelay(); + + if ( + ProgressState != ProgressState.Inactive + && ProgressState != ProgressState.Paused + && ProgressState != ProgressState.Pending + ) { Logger.Warn( "Attempted to resume download {Download} but it is not paused ({State})", FileName, ProgressState ); + return; } Logger.Debug("Resuming download {Download}", FileName); @@ -235,6 +256,9 @@ internal void Resume() public void Pause() { + // Cancel any pending auto-retry delay. + CancelRetryDelay(); + if (ProgressState != ProgressState.Working) { Logger.Warn( @@ -252,6 +276,33 @@ public void Pause() OnProgressStateChanged(ProgressState); } + /// + /// Cleans up temp file and all sidecar files (e.g. .cm-info.json, preview image) + /// for a download that has already failed and will not be retried. + /// This transitions the state to so the + /// service removes the tracking entry. + /// + public void Dismiss() + { + if (ProgressState != ProgressState.Failed) + { + Logger.Warn( + "Attempted to dismiss download {Download} but it is not in a failed state ({State})", + FileName, + ProgressState + ); + return; + } + + Logger.Debug("Dismissing failed download {Download}", FileName); + + DoCleanup(); + + OnProgressStateChanging(ProgressState.Cancelled); + ProgressState = ProgressState.Cancelled; + OnProgressStateChanged(ProgressState); + } + public void Cancel() { if (ProgressState is not (ProgressState.Working or ProgressState.Inactive)) @@ -264,6 +315,9 @@ public void Cancel() return; } + // Cancel any pending auto-retry delay. + CancelRetryDelay(); + Logger.Debug("Cancelling download {Download}", FileName); // Cancel token if it exists @@ -290,9 +344,12 @@ public void SetPending() } /// - /// Deletes the temp file and any extra cleanup files + /// Deletes the temp file and, optionally, any extra cleanup files (e.g. sidecar metadata). + /// Pass as false when the download + /// failed but may be retried — sidecar files (.cm-info.json, preview image) should survive + /// so a manual retry doesn't need to recreate them. /// - private void DoCleanup() + private void DoCleanup(bool includeExtraCleanupFiles = true) { try { @@ -303,6 +360,9 @@ private void DoCleanup() Logger.Warn("Failed to delete temp file {TempFile}", TempFileName); } + if (!includeExtraCleanupFiles) + return; + foreach (var extraFile in ExtraCleanupFileNames) { try @@ -316,6 +376,16 @@ private void DoCleanup() } } + /// + /// Returns true for transient network/SSL exceptions that are safe to retry (ie: VPN tunnel resets or TLS re-key failures) + /// (IOException, AuthenticationException, or either wrapped in an AggregateException). + /// + private static bool IsTransientNetworkException(Exception? ex) => + ex is IOException or AuthenticationException + || ex?.InnerException is IOException or AuthenticationException + || ex is AggregateException ae + && ae.InnerExceptions.Any(e => e is IOException or AuthenticationException); + /// /// Invoked by the task's completion callback /// @@ -349,7 +419,7 @@ private void OnDownloadTaskCompleted(Task task) // Set the exception Exception = task.Exception; - if ((Exception is IOException || Exception?.InnerException is IOException) && attempts < 3) + if (IsTransientNetworkException(Exception) && attempts < MaxRetryAttempts) { attempts++; Logger.Warn( @@ -359,9 +429,39 @@ private void OnDownloadTaskCompleted(Task task) attempts ); + // Exponential backoff: 2 s → 4 s → 8 s, capped at 30 s, ±500 ms jitter. + // Gives the VPN tunnel time to re-key/re-route before reconnecting, + // which prevents the retry from hitting the same torn connection. + var delayMs = + (int)Math.Min(2000 * Math.Pow(2, attempts - 1), 30_000) + Random.Shared.Next(-500, 500); + Logger.Debug( + "Download {Download} retrying in {Delay}ms (attempt {Attempt}/{MaxAttempts})", + FileName, + delayMs, + attempts, + MaxRetryAttempts + ); + + // Persist Inactive to disk before the delay so a restart during backoff loads it as resumable. OnProgressStateChanging(ProgressState.Inactive); ProgressState = ProgressState.Inactive; - Resume(); + OnProgressStateChanged(ProgressState.Inactive); + + // Clean up the completed task resources; Resume() will create new ones. + downloadTask = null; + downloadCancellationTokenSource = null; + downloadPauseTokenSource = null; + + // Schedule the retry with a cancellation token so Cancel/Pause can abort the delay. + retryDelayCancellationTokenSource?.Dispose(); + retryDelayCancellationTokenSource = new CancellationTokenSource(); + Task.Delay(Math.Max(delayMs, 0), retryDelayCancellationTokenSource.Token) + .ContinueWith(t => + { + if (t.IsCompletedSuccessfully) + Resume(); + }) + .SafeFireAndForget(); return; } @@ -377,11 +477,17 @@ private void OnDownloadTaskCompleted(Task task) ProgressState = ProgressState.Success; } - // For failed or cancelled, delete the temp files - if (ProgressState is ProgressState.Failed or ProgressState.Cancelled) + // For cancelled, delete the temp file and any sidecar metadata. + // For failed, only delete the temp file — sidecar files (.cm-info.json, preview image) + // are preserved so a manual retry doesn't need to recreate them. + if (ProgressState is ProgressState.Cancelled) { DoCleanup(); } + else if (ProgressState is ProgressState.Failed) + { + DoCleanup(includeExtraCleanupFiles: false); + } // For pause, just do nothing OnProgressStateChanged(ProgressState); @@ -392,6 +498,17 @@ private void OnDownloadTaskCompleted(Task task) downloadPauseTokenSource = null; } + /// + /// Resets the retry counter and silently sets state to Inactive without firing events. + /// Must be called before re-adding to TrackedDownloadService to avoid events + /// firing while the download is absent from the dictionary. + /// + public void ResetAttempts() + { + attempts = 0; + ProgressState = ProgressState.Inactive; + } + public void SetDownloadService(IDownloadService service) { downloadService = service; diff --git a/StabilityMatrix.Core/Services/ITrackedDownloadService.cs b/StabilityMatrix.Core/Services/ITrackedDownloadService.cs index ee1e2ba8..86c00da3 100644 --- a/StabilityMatrix.Core/Services/ITrackedDownloadService.cs +++ b/StabilityMatrix.Core/Services/ITrackedDownloadService.cs @@ -15,5 +15,7 @@ TrackedDownload NewDownload(string downloadUrl, FilePath downloadPath) => NewDownload(new Uri(downloadUrl), downloadPath); Task TryStartDownload(TrackedDownload download); Task TryResumeDownload(TrackedDownload download); + Task TryRestartDownload(TrackedDownload download); + void UpdateMaxConcurrentDownloads(int newMax); } diff --git a/StabilityMatrix.Core/Services/TrackedDownloadService.cs b/StabilityMatrix.Core/Services/TrackedDownloadService.cs index 12cf3ca7..3db919db 100644 --- a/StabilityMatrix.Core/Services/TrackedDownloadService.cs +++ b/StabilityMatrix.Core/Services/TrackedDownloadService.cs @@ -129,6 +129,46 @@ public async Task TryStartDownload(TrackedDownload download) } } + public async Task TryRestartDownload(TrackedDownload download) + { + // Re-create the backing JSON file and re-add to the dictionary. + // Downloads are removed on failure, so this restores the tracking entry + // so that subsequent state-change events can persist normally. + var downloadsDir = new DirectoryPath(settingsManager.DownloadsDirectory); + downloadsDir.Create(); + var jsonFile = downloadsDir.JoinFile($"{download.Id}.json"); + + var jsonFileStream = new FileStream( + jsonFile.Info.FullName, + FileMode.Create, + FileAccess.ReadWrite, + FileShare.Read, + bufferSize: 4096, + useAsync: true + ); + var jsonBytes = JsonSerializer.SerializeToUtf8Bytes(download); + + try + { + await jsonFileStream.WriteAsync(jsonBytes).ConfigureAwait(false); + await jsonFileStream.FlushAsync().ConfigureAwait(false); + + // Handlers are already attached from the original AddDownload call. + if (!downloads.TryAdd(download.Id, (download, jsonFileStream))) + { + // Already tracked; discard the newly opened stream. + await jsonFileStream.DisposeAsync().ConfigureAwait(false); + } + } + catch + { + await jsonFileStream.DisposeAsync().ConfigureAwait(false); + throw; + } + + await TryResumeDownload(download).ConfigureAwait(false); + } + public async Task TryResumeDownload(TrackedDownload download) { if (IsQueueEnabled && ActiveDownloads >= MaxConcurrentDownloads) @@ -205,12 +245,17 @@ private void AdjustSemaphoreCount() /// private void UpdateJsonForDownload(TrackedDownload download) { + // The download may have already been removed from the dictionary (e.g. it failed and was + // then dismissed). In that case there is nothing to persist, so skip silently. + if (!downloads.TryGetValue(download.Id, out var downloadInfo)) + return; + // Serialize to json var json = JsonSerializer.Serialize(download); var jsonBytes = Encoding.UTF8.GetBytes(json); // Write to file - var (_, fs) = downloads[download.Id]; + var (_, fs) = downloadInfo; fs.Seek(0, SeekOrigin.Begin); fs.Write(jsonBytes); fs.Flush(); diff --git a/StabilityMatrix.Tests/Models/TrackedDownloadTests.cs b/StabilityMatrix.Tests/Models/TrackedDownloadTests.cs new file mode 100644 index 00000000..63804534 --- /dev/null +++ b/StabilityMatrix.Tests/Models/TrackedDownloadTests.cs @@ -0,0 +1,178 @@ +using System; +using System.IO; +using System.Threading; +using System.Threading.Tasks; +using NSubstitute; +using StabilityMatrix.Core.Models; +using StabilityMatrix.Core.Models.FileInterfaces; +using StabilityMatrix.Core.Models.Progress; +using StabilityMatrix.Core.Services; + +namespace StabilityMatrix.Tests.Models; + +[TestClass] +public class TrackedDownloadTests +{ + private string tempDir = null!; + + [TestInitialize] + public void Setup() + { + tempDir = Path.Combine(Path.GetTempPath(), Path.GetRandomFileName()); + Directory.CreateDirectory(tempDir); + } + + [TestCleanup] + public void Cleanup() + { + if (Directory.Exists(tempDir)) + Directory.Delete(tempDir, recursive: true); + } + + private TrackedDownload CreateDownload(IDownloadService downloadService) + { + var download = new TrackedDownload + { + Id = Guid.NewGuid(), + SourceUrl = new Uri("https://example.com/model.safetensors"), + DownloadDirectory = new DirectoryPath(tempDir), + FileName = "model.safetensors", + TempFileName = "model.safetensors.partial", + }; + download.SetDownloadService(downloadService); + return download; + } + + // Resume() must proceed when the state is Pending (queued resume fix). + + [TestMethod] + public async Task Resume_WhileInPendingState_SetsStateToWorking() + { + // Arrange – download service that blocks forever until cancelled. + var downloadService = Substitute.For(); + downloadService + .ResumeDownloadToFileAsync( + Arg.Any(), + Arg.Any(), + Arg.Any(), + Arg.Any>(), + Arg.Any(), + Arg.Any() + ) + .Returns(callInfo => + { + var ct = callInfo.Arg(); + return Task.Delay(Timeout.Infinite, ct); + }); + + var download = CreateDownload(downloadService); + download.SetPending(); // Simulate being queued + + // Act + download.Resume(); + + // Assert – Resume() must not have returned early; state is now Working. + Assert.AreEqual(ProgressState.Working, download.ProgressState); + + // Cleanup – cancel and wait for the task to finish. + var cancelledTcs = new TaskCompletionSource(); + download.ProgressStateChanged += (_, state) => + { + if (state == ProgressState.Cancelled) + cancelledTcs.TrySetResult(); + }; + download.Cancel(); + await cancelledTcs.Task.WaitAsync(TimeSpan.FromSeconds(5)); + } + + // Sidecar files (.cm-info.json, preview image) must survive a failed + // download so a manual retry can succeed without recreating them. + + [TestMethod] + public async Task OnFailed_SidecarFilesPreservedForRetry() + { + // Arrange – create sidecar files that ModelImportService would have written. + var sidecarPath = Path.Combine(tempDir, "model.cm-info.json"); + var previewPath = Path.Combine(tempDir, "model.preview.png"); + await File.WriteAllTextAsync(sidecarPath, "{}"); + await File.WriteAllTextAsync(previewPath, "PNG"); + + // Download service that immediately throws a non-transient error + // (InvalidOperationException is not an IOException/AuthenticationException, + // so it goes straight to Failed without any auto-retry attempts). + var downloadService = Substitute.For(); + downloadService + .ResumeDownloadToFileAsync( + Arg.Any(), + Arg.Any(), + Arg.Any(), + Arg.Any>(), + Arg.Any(), + Arg.Any() + ) + .Returns(Task.FromException(new InvalidOperationException("Simulated download error"))); + + var download = CreateDownload(downloadService); + download.ExtraCleanupFileNames.Add(sidecarPath); + download.ExtraCleanupFileNames.Add(previewPath); + + var failedTcs = new TaskCompletionSource(); + download.ProgressStateChanged += (_, state) => + { + if (state == ProgressState.Failed) + failedTcs.TrySetResult(); + }; + + // Act + download.Start(); + await failedTcs.Task.WaitAsync(TimeSpan.FromSeconds(5)); + + // Assert – sidecar files must still exist for a potential manual retry. + Assert.IsTrue(File.Exists(sidecarPath), ".cm-info.json should be preserved after failure"); + Assert.IsTrue(File.Exists(previewPath), "Preview image should be preserved after failure"); + } + + [TestMethod] + public async Task OnCancelled_SidecarFilesAreDeleted() + { + // Sidecar files should still be cleaned up when the user explicitly cancels. + var sidecarPath = Path.Combine(tempDir, "model.cm-info.json"); + var previewPath = Path.Combine(tempDir, "model.preview.png"); + await File.WriteAllTextAsync(sidecarPath, "{}"); + await File.WriteAllTextAsync(previewPath, "PNG"); + + var downloadService = Substitute.For(); + downloadService + .ResumeDownloadToFileAsync( + Arg.Any(), + Arg.Any(), + Arg.Any(), + Arg.Any>(), + Arg.Any(), + Arg.Any() + ) + .Returns(callInfo => + { + var ct = callInfo.Arg(); + return Task.Delay(Timeout.Infinite, ct); + }); + + var download = CreateDownload(downloadService); + download.ExtraCleanupFileNames.Add(sidecarPath); + download.ExtraCleanupFileNames.Add(previewPath); + + var cancelledTcs = new TaskCompletionSource(); + download.ProgressStateChanged += (_, state) => + { + if (state == ProgressState.Cancelled) + cancelledTcs.TrySetResult(); + }; + + download.Start(); + download.Cancel(); + await cancelledTcs.Task.WaitAsync(TimeSpan.FromSeconds(5)); + + Assert.IsFalse(File.Exists(sidecarPath), ".cm-info.json should be deleted on cancel"); + Assert.IsFalse(File.Exists(previewPath), "Preview image should be deleted on cancel"); + } +}