Merge pull request LykosAI#1574 from NeuralFault/vpn-ssl-downloadfix

Fix: Model downloads failing on VPN connections (SSL/TLS decrypt error). QoL: Retry button

Author: mohnjiles

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;
 
+    /// <summary>
+    /// Override to true in subclasses that support manual retry after failure.
+    /// Defaults to false so unrelated progress item types are never affected.
+    /// </summary>
+    public virtual bool SupportsRetry => false;
+
+    /// <summary>
+    /// Override to true in subclasses that support dismissing a failed item,
+    /// which runs full sidecar cleanup before removing the entry.
+    /// </summary>
+    public virtual bool SupportsDismiss => false;
+
     public bool CanPauseResume => SupportsPauseResume && !IsCompleted && !IsPending;
     public bool CanCancel => SupportsCancel && !IsCompleted;
 
+    /// <summary>
+    /// True only when this item supports retry AND is in the Failed state.
+    /// </summary>
+    public bool CanRetry => SupportsRetry && State == ProgressState.Failed;
+
+    /// <summary>
+    /// True only when this item supports dismiss AND is in the Failed state.
+    /// </summary>
+    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()
     {

StabilityMatrix.Avalonia/ViewModels/Progress/DownloadProgressItemViewModel.cs

@@ -71,6 +71,17 @@ private void OnProgressStateChanged(ProgressState state)
         }
     }
 
+    /// <summary>
+    /// Downloads support manual retry when they reach the Failed state.
+    /// </summary>
+    public override bool SupportsRetry => true;
+
+    /// <summary>
+    /// Downloads support dismiss, which cleans up all sidecar files when
+    /// the user discards a failed download without retrying.
+    /// </summary>
+    public override bool SupportsDismiss => true;
+
     /// <inheritdoc />
     public override Task Cancel()
     {
@@ -91,4 +102,23 @@ public override Task Resume()
     {
         return downloadService.TryResumeDownload(download);
     }
+
+    /// <inheritdoc />
+    /// 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);
+    }
+
+    /// <inheritdoc />
+    /// 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;
+    }
 }

StabilityMatrix.Avalonia/Views/ProgressManagerPage.axaml

@@ -113,6 +113,24 @@
                                         IsVisible="{Binding CanCancel}">
                                         <ui:SymbolIcon Symbol="Cancel" />
                                     </Button>
+
+                                    <!--  Retry button: only visible when download is in Failed state  -->
+                                    <Button
+                                        Classes="transparent-full"
+                                        Command="{Binding RetryCommand}"
+                                        IsVisible="{Binding CanRetry}"
+                                        ToolTip.Tip="Retry download">
+                                        <ui:SymbolIcon Symbol="Refresh" />
+                                    </Button>
+
+                                    <!--  Dismiss button: cleans up sidecar files for failed downloads that won't be retried  -->
+                                    <Button
+                                        Classes="transparent-full"
+                                        Command="{Binding DismissCommand}"
+                                        IsVisible="{Binding CanDismiss}"
+                                        ToolTip.Tip="Dismiss and clean up">
+                                        <ui:SymbolIcon Symbol="Delete" />
+                                    </Button>
                                 </StackPanel>
 
                                 <ProgressBar

StabilityMatrix.Core/Models/TrackedDownload.cs

@@ -1,4 +1,5 @@
 using System.Diagnostics.CodeAnalysis;
+using System.Security.Authentication;
 using System.Text.Json.Serialization;
 using AsyncAwaitBestPractices;
 using NLog;
@@ -77,7 +78,9 @@ public class TrackedDownload
     [JsonIgnore]
     public Exception? Exception { get; private set; }
 
+    private const int MaxRetryAttempts = 3;
     private int attempts;
+    private CancellationTokenSource? retryDelayCancellationTokenSource;
 
     #region Events
     public event EventHandler<ProgressReport>? 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<ProgressReport>(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);
     }
 
+    /// <summary>
+    /// 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 <see cref="ProgressState.Cancelled"/> so the
+    /// service removes the tracking entry.
+    /// </summary>
+    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()
     }
 
     /// <summary>
-    /// Deletes the temp file and any extra cleanup files
+    /// Deletes the temp file and, optionally, any extra cleanup files (e.g. sidecar metadata).
+    /// Pass <paramref name="includeExtraCleanupFiles"/> as <c>false</c> 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.
     /// </summary>
-    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()
         }
     }
 
+    /// <summary>
+    /// 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).
+    /// </summary>
+    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);
+
     /// <summary>
     /// Invoked by the task's completion callback
     /// </summary>
@@ -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;
     }
 
+    /// <summary>
+    /// 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.
+    /// </summary>
+    public void ResetAttempts()
+    {
+        attempts = 0;
+        ProgressState = ProgressState.Inactive;
+    }
+
     public void SetDownloadService(IDownloadService service)
     {
         downloadService = service;

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