From 107a983868a54ea817dc4196d0a7e2a2ec2f1c15 Mon Sep 17 00:00:00 2001
From: Ionite <dev@ionite.io>
Date: Fri, 16 Jun 2023 17:10:11 -0400
Subject: [PATCH] Add cancel token to A3WebApi

---
 StabilityMatrix/Api/IA3WebApi.cs              | 13 ++++----
 .../ViewModels/SettingsViewModel.cs           | 33 +++++++------------
 2 files changed, 18 insertions(+), 28 deletions(-)

diff --git a/StabilityMatrix/Api/IA3WebApi.cs b/StabilityMatrix/Api/IA3WebApi.cs
index 71c1d48b..506d351a 100644
--- a/StabilityMatrix/Api/IA3WebApi.cs
+++ b/StabilityMatrix/Api/IA3WebApi.cs
@@ -1,4 +1,5 @@
-using System.Threading.Tasks;
+using System.Threading;
+using System.Threading.Tasks;
 using Refit;
 using StabilityMatrix.Models.Api;
 
@@ -8,17 +9,17 @@ namespace StabilityMatrix.Api;
 public interface IA3WebApi
 {
     [Get("/internal/ping")]
-    Task<string> GetPing();
+    Task<string> GetPing(CancellationToken cancellationToken = default);
     
     [Post("/sdapi/v1/txt2img")]
-    Task<ImageResponse> TextToImage([Body] TextToImageRequest request);
+    Task<ImageResponse> TextToImage([Body] TextToImageRequest request, CancellationToken cancellationToken = default);
     
     [Get("/sdapi/v1/progress")]
-    Task<ProgressResponse> GetProgress([Body] ProgressRequest request);
+    Task<ProgressResponse> GetProgress([Body] ProgressRequest request, CancellationToken cancellationToken = default);
     
     [Get("/sdapi/v1/options")]
-    Task<A3Options> GetOptions();
+    Task<A3Options> GetOptions(CancellationToken cancellationToken = default);
     
     [Post("/sdapi/v1/options")]
-    Task SetOptions([Body] A3Options request);
+    Task SetOptions([Body] A3Options request, CancellationToken cancellationToken = default);
 }
diff --git a/StabilityMatrix/ViewModels/SettingsViewModel.cs b/StabilityMatrix/ViewModels/SettingsViewModel.cs
index b392c828..97cfa8cd 100644
--- a/StabilityMatrix/ViewModels/SettingsViewModel.cs
+++ b/StabilityMatrix/ViewModels/SettingsViewModel.cs
@@ -8,6 +8,7 @@ using System.IO;
 using System.Linq;
 using System.Reflection;
 using System.Text.Json;
+using System.Threading;
 using System.Threading.Tasks;
 using System.Windows;
 using System.Windows.Controls;
@@ -18,6 +19,7 @@ using CommunityToolkit.Mvvm.Input;
 using MdXaml;
 using Microsoft.Extensions.Logging;
 using Ookii.Dialogs.Wpf;
+using Polly.Timeout;
 using Refit;
 using StabilityMatrix.Api;
 using StabilityMatrix.Helper;
@@ -216,37 +218,24 @@ public partial class SettingsViewModel : ObservableObject
         }
     }
 
-    private Task<bool> TryPingWebApi()
+    private async Task<bool> TryPingWebApi()
     {
-        return TryPingWebApi(TimeSpan.FromMilliseconds(150));
-    }
-
-    private async Task<bool> TryPingWebApi(TimeSpan? minimumDelay)
-    {
-        bool result;
-        var timer = minimumDelay != null ? Stopwatch.StartNew() : null;
+        await using var minimumDelay = new MinimumDelay(100, 200);
         try
         {
             await a3WebApiManager.Client.GetPing();
-            result = true;
+            return true;
         }
-        catch (Exception ex)
+        catch (TimeoutRejectedException)
         {
-            logger.LogInformation(ex, "Ping failed");
-            result = false;
+            logger.LogInformation("Ping timed out");
+            return false;
         }
-        
-        // Wait if minimum delay is set
-        if (minimumDelay != null)
+        catch (ApiException ex)
         {
-            var delay = minimumDelay.Value - timer!.Elapsed;
-            if (delay > TimeSpan.Zero)
-            {
-                await Task.Delay(delay);
-            }
+            logger.LogInformation("Ping failed with status [{StatusCode}]: {Content}", ex.StatusCode, ex.ReasonPhrase);
+            return false;
         }
-
-        return result;
     }
     
     [RelayCommand]