From c67f17cf41e34f44762f43806f066373db911ad8 Mon Sep 17 00:00:00 2001
From: Thorsten Sommer <SommerEngineering@users.noreply.github.com>
Date: Mon, 17 Mar 2025 18:00:26 +0100
Subject: [PATCH] Added a "stop generation" button to assistants (#365)

---
 .../Assistants/AssistantBase.razor            | 15 ++++++++----
 .../Assistants/AssistantBase.razor.cs         | 24 ++++++++++++++-----
 .../wwwroot/changelog/v0.9.38.md              |  1 +
 3 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/app/MindWork AI Studio/Assistants/AssistantBase.razor b/app/MindWork AI Studio/Assistants/AssistantBase.razor
index f5883156..38d4b8b8 100644
--- a/app/MindWork AI Studio/Assistants/AssistantBase.razor	
+++ b/app/MindWork AI Studio/Assistants/AssistantBase.razor	
@@ -28,10 +28,17 @@
                         @this.Body
                     </CascadingValue>
                     
-                    <MudButton Disabled="@this.SubmitDisabled" Variant="Variant.Filled" Class="mb-3" OnClick="() => this.SubmitAction()" Style="@this.SubmitButtonStyle">
-                        @this.SubmitText
-                    </MudButton>
-                    
+                    <MudStack Row="true" AlignItems="AlignItems.Center" StretchItems="StretchItems.Start" Class="mb-3">
+                        <MudButton Disabled="@this.SubmitDisabled" Variant="Variant.Filled" OnClick="() => this.SubmitAction()" Style="@this.SubmitButtonStyle">
+                            @this.SubmitText
+                        </MudButton>
+                        @if (this.isProcessing && this.cancellationTokenSource is not null)
+                        {
+                            <MudTooltip Text="Stop generation">
+                                <MudIconButton Variant="Variant.Filled" Icon="@Icons.Material.Filled.Stop" Color="Color.Error" OnClick="() => this.CancelStreaming()"/>
+                            </MudTooltip>
+                        }
+                    </MudStack>
                 }
             </MudForm>
             <Issues IssuesData="@(this.inputIssues)"/>
diff --git a/app/MindWork AI Studio/Assistants/AssistantBase.razor.cs b/app/MindWork AI Studio/Assistants/AssistantBase.razor.cs
index c1d5b8e8..1f0c7364 100644
--- a/app/MindWork AI Studio/Assistants/AssistantBase.razor.cs	
+++ b/app/MindWork AI Studio/Assistants/AssistantBase.razor.cs	
@@ -100,6 +100,7 @@ public abstract partial class AssistantBase<TSettings> : AssistantLowerBase, IMe
     
     private readonly Timer formChangeTimer = new(TimeSpan.FromSeconds(1.6));
     
+    private CancellationTokenSource? cancellationTokenSource;
     private ContentBlock? resultingContentBlock;
     private string[] inputIssues = [];
     private bool isProcessing;
@@ -285,12 +286,16 @@ public abstract partial class AssistantBase<TSettings> : AssistantLowerBase, IMe
 
         this.isProcessing = true;
         this.StateHasChanged();
-        
-        // Use the selected provider to get the AI response.
-        // By awaiting this line, we wait for the entire
-        // content to be streamed.
-        this.chatThread = await aiText.CreateFromProviderAsync(this.providerSettings.CreateProvider(this.Logger), this.providerSettings.Model, this.lastUserPrompt, this.chatThread);
-        
+
+        using (this.cancellationTokenSource = new())
+        {
+            // Use the selected provider to get the AI response.
+            // By awaiting this line, we wait for the entire
+            // content to be streamed.
+            this.chatThread = await aiText.CreateFromProviderAsync(this.providerSettings.CreateProvider(this.Logger), this.providerSettings.Model, this.lastUserPrompt, this.chatThread, this.cancellationTokenSource.Token);
+        }
+
+        this.cancellationTokenSource = null;
         this.isProcessing = false;
         this.StateHasChanged();
         
@@ -298,6 +303,13 @@ public abstract partial class AssistantBase<TSettings> : AssistantLowerBase, IMe
         return aiText.Text;
     }
     
+    private async Task CancelStreaming()
+    {
+        if (this.cancellationTokenSource is not null)
+            if(!this.cancellationTokenSource.IsCancellationRequested)
+                await this.cancellationTokenSource.CancelAsync();
+    }
+    
     protected async Task CopyToClipboard()
     {
         await this.RustService.CopyText2Clipboard(this.Snackbar, this.Result2Copy());
diff --git a/app/MindWork AI Studio/wwwroot/changelog/v0.9.38.md b/app/MindWork AI Studio/wwwroot/changelog/v0.9.38.md
index 86682570..00144ef4 100644
--- a/app/MindWork AI Studio/wwwroot/changelog/v0.9.38.md	
+++ b/app/MindWork AI Studio/wwwroot/changelog/v0.9.38.md	
@@ -1,2 +1,3 @@
 # v0.9.38, build 213 (2025-03-xx xx:xx UTC)
+- Added the "stop generation" button to all assistants.
 - Updated the ERI v1 specification for the ERI server assistant & fixed spelling of the `UNKNOWN` role.
\ No newline at end of file