Fixed issues with parallel chats and switching chat threads

app/MindWork AI Studio/Components/ChatComponent.razor.cs

@@ -70,6 +70,8 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
     private string currentWorkspaceName = string.Empty;
     private Guid currentWorkspaceId = Guid.Empty;
     private Guid currentChatThreadId = Guid.Empty;
+    private Guid loadedParameterChatId = Guid.Empty;
+    private Guid loadedParameterWorkspaceId = Guid.Empty;
     private Guid foregroundChatId = Guid.Empty;
     private int workspaceHeaderSyncVersion;
     private HashSet<FileAttachment> chatDocumentPaths = [];
@@ -83,7 +85,7 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
     protected override async Task OnInitializedAsync()
     {
         // Apply the filters for the message bus:
-        this.ApplyFilters([], [ Event.HAS_CHAT_UNSAVED_CHANGES, Event.RESET_CHAT_STATE, Event.CHAT_STREAMING_DONE, Event.WORKSPACE_LOADED_CHAT_CHANGED, Event.AI_JOB_CHANGED, Event.AI_JOB_FINISHED, Event.CHAT_GENERATION_CHANGED ]);
+        this.ApplyFilters([], [ Event.HAS_CHAT_UNSAVED_CHANGES, Event.RESET_CHAT_STATE, Event.CHAT_STREAMING_DONE, Event.AI_JOB_CHANGED, Event.AI_JOB_FINISHED, Event.CHAT_GENERATION_CHANGED ]);
         
         // Configure the spellchecking for the user input:
         this.SettingsManager.InjectSpellchecking(USER_INPUT_ATTRIBUTES);
@@ -120,6 +122,7 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
             this.ChatThread.IncludeDateTime = true;
             
             this.Logger.LogInformation($"The chat '{this.ChatThread.ChatId}' with {this.ChatThread.Blocks.Count} messages was deferred and will be rendered now.");
+            this.MarkCurrentChatAsLoadedParameter();
             await this.ChatThreadChanged.InvokeAsync(this.ChatThread);
             
             // We know already that the chat thread is not null,
@@ -246,6 +249,7 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
             
             if(this.ChatThread is not null)
             {
+                this.MarkCurrentChatAsLoadedParameter();
                 await this.ChatThreadChanged.InvokeAsync(this.ChatThread);
                 this.Logger.LogInformation($"The chat '{this.ChatThread!.ChatId}' with title '{this.ChatThread.Name}' ({this.ChatThread.Blocks.Count} messages) was loaded successfully.");
 
@@ -276,13 +280,35 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
 
     protected override async Task OnParametersSetAsync()
     {
-        await this.SyncWorkspaceHeaderWithChatThreadAsync();
+        await this.ApplyLoadedChatParameterAsync();
         await this.SyncForegroundChatAsync();
         await base.OnParametersSetAsync();
     }
 
     #endregion
 
+    private async Task ApplyLoadedChatParameterAsync()
+    {
+        var chatId = this.ChatThread?.ChatId ?? Guid.Empty;
+        var workspaceId = this.ChatThread?.WorkspaceId ?? Guid.Empty;
+
+        if (this.loadedParameterChatId == chatId && this.loadedParameterWorkspaceId == workspaceId)
+        {
+            await this.SyncWorkspaceHeaderWithChatThreadAsync();
+            return;
+        }
+
+        this.loadedParameterChatId = chatId;
+        this.loadedParameterWorkspaceId = workspaceId;
+        await this.LoadedChatChanged(notifyParent: false);
+    }
+
+    private void MarkCurrentChatAsLoadedParameter()
+    {
+        this.loadedParameterChatId = this.ChatThread?.ChatId ?? Guid.Empty;
+        this.loadedParameterWorkspaceId = this.ChatThread?.WorkspaceId ?? Guid.Empty;
+    }
+
     private async Task SyncWorkspaceHeaderWithChatThreadAsync()
     {
         var syncVersion = Interlocked.Increment(ref this.workspaceHeaderSyncVersion);
@@ -552,6 +578,7 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
                 Blocks = this.currentChatTemplate == ChatTemplate.NO_CHAT_TEMPLATE ? [] : this.currentChatTemplate.ExampleConversation.Select(x => x.DeepClone()).ToList(),
             };
             
+            this.MarkCurrentChatAsLoadedParameter();
             await this.ChatThreadChanged.InvokeAsync(this.ChatThread);
         }
         else
@@ -793,6 +820,7 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
         
         // Notify the parent component about the change:
         await this.SyncForegroundChatAsync();
+        this.MarkCurrentChatAsLoadedParameter();
         await this.ChatThreadChanged.InvokeAsync(this.ChatThread);
     }
     
@@ -834,12 +862,13 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
         await WorkspaceBehaviour.DeleteChatAsync(this.DialogService, this.ChatThread!.WorkspaceId, this.ChatThread.ChatId, askForConfirmation: false);
         
         this.ChatThread!.WorkspaceId = workspaceId;
+        this.MarkCurrentChatAsLoadedParameter();
         await this.SaveThread();
 
         await this.SyncWorkspaceHeaderWithChatThreadAsync();
     }
     
-    private async Task LoadedChatChanged()
+    private async Task LoadedChatChanged(bool notifyParent = true)
     {
         this.hasUnsavedChanges = false;
         this.userInput = string.Empty;
@@ -847,13 +876,19 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
         if (this.ChatThread is not null)
         {
             this.ChatThread = this.AIJobService.TryGetLiveChatThread(this.ChatThread.ChatId) ?? this.ChatThread;
-            await this.ChatThreadChanged.InvokeAsync(this.ChatThread);
+            this.loadedParameterChatId = this.ChatThread.ChatId;
+            this.loadedParameterWorkspaceId = this.ChatThread.WorkspaceId;
+            if (notifyParent)
+                await this.ChatThreadChanged.InvokeAsync(this.ChatThread);
+
             await this.SyncWorkspaceHeaderWithChatThreadAsync();
             await this.SyncForegroundChatAsync();
             this.dataSourceSelectionComponent?.ChangeOptionWithoutSaving(this.ChatThread.DataSourceOptions, this.ChatThread.AISelectedDataSources);
         }
         else
         {
+            this.loadedParameterChatId = Guid.Empty;
+            this.loadedParameterWorkspaceId = Guid.Empty;
             this.ClearWorkspaceHeaderState();
             await this.SyncForegroundChatAsync();
             this.ApplyStandardDataSourceOptions();
@@ -876,6 +911,7 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
         this.ClearWorkspaceHeaderState();
         
         this.ChatThread = null;
+        this.MarkCurrentChatAsLoadedParameter();
         await this.SyncForegroundChatAsync();
         this.ApplyStandardDataSourceOptions();
         await this.ChatThreadChanged.InvokeAsync(this.ChatThread);
@@ -1000,10 +1036,6 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
                     await this.SaveThread();
                 break;
             
-            case Event.WORKSPACE_LOADED_CHAT_CHANGED:
-                await this.LoadedChatChanged();
-                break;
-
             case Event.AI_JOB_CHANGED:
             case Event.AI_JOB_FINISHED:
             case Event.CHAT_GENERATION_CHANGED:
@@ -1049,6 +1081,7 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
         }
 
         await this.AIJobService.SetForegroundAsync(AIJobKind.CHAT_GENERATION, this.foregroundChatId, false);
+        this.Dispose();
     }
 
     #endregion

app/MindWork AI Studio/Components/Workspaces.razor.cs

@@ -236,12 +236,13 @@ public partial class Workspaces : MSGComponentBase
     private string GetChatTreeIcon(Guid chatId, string defaultIcon)
     {
         var snapshot = this.AIJobService.TryGetChatSnapshot(chatId);
-        return snapshot?.Status switch
+        if (snapshot is null || !snapshot.IsActive)
+            return defaultIcon;
+
+        return snapshot.Status switch
         {
             AIJobStatus.WAITING_FOR_REMOTE => Icons.Material.Filled.HourglassTop,
             AIJobStatus.RUNNING => Icons.Material.Filled.ChangeCircle,
-            AIJobStatus.CANCELED => Icons.Material.Filled.Cancel,
-            AIJobStatus.FAILED => Icons.Material.Filled.Error,
             _ => defaultIcon,
         };
     }
@@ -390,7 +391,6 @@ public partial class Workspaces : MSGComponentBase
             {
                 this.CurrentChatThread = chat;
                 await this.CurrentChatThreadChanged.InvokeAsync(this.CurrentChatThread);
-                await MessageBus.INSTANCE.SendMessage<bool>(this, Event.WORKSPACE_LOADED_CHAT_CHANGED);
             }
             
             return chat;
@@ -439,7 +439,6 @@ public partial class Workspaces : MSGComponentBase
         {
             this.CurrentChatThread = null;
             await this.CurrentChatThreadChanged.InvokeAsync(this.CurrentChatThread);
-            await MessageBus.INSTANCE.SendMessage<bool>(this, Event.WORKSPACE_LOADED_CHAT_CHANGED);
         }
     }
 
@@ -473,7 +472,6 @@ public partial class Workspaces : MSGComponentBase
         {
             this.CurrentChatThread.Name = chat.Name;
             await this.CurrentChatThreadChanged.InvokeAsync(this.CurrentChatThread);
-            await MessageBus.INSTANCE.SendMessage<bool>(this, Event.WORKSPACE_LOADED_CHAT_CHANGED);
         }
         
         await WorkspaceBehaviour.StoreChatAsync(chat);
@@ -596,7 +594,6 @@ public partial class Workspaces : MSGComponentBase
         {
             this.CurrentChatThread = chat;
             await this.CurrentChatThreadChanged.InvokeAsync(this.CurrentChatThread);
-            await MessageBus.INSTANCE.SendMessage<bool>(this, Event.WORKSPACE_LOADED_CHAT_CHANGED);
         }
         
         await WorkspaceBehaviour.StoreChatAsync(chat);

app/MindWork AI Studio/Tools/AIJobs/AIJobService.cs

@@ -17,12 +17,16 @@ public sealed class AIJobService(
     {
         public required CancellationTokenSource CancellationTokenSource { get; init; }
 
+        public required CancellationToken CancellationToken { get; init; }
+
         public required ChatGenerationRequest ChatGenerationRequest { get; init; }
 
         public required AIJobSnapshot Snapshot { get; set; }
 
         public DateTimeOffset LastCheckpoint { get; set; }
 
+        public bool IsCompletionStarted { get; set; }
+
         public readonly Lock SyncRoot = new();
     }
 
@@ -96,9 +100,11 @@ public sealed class AIJobService(
             UpdatedAt = DateTimeOffset.Now,
         };
 
+        var cancellationTokenSource = new CancellationTokenSource();
         var state = new AIJobState
         {
-            CancellationTokenSource = new CancellationTokenSource(),
+            CancellationTokenSource = cancellationTokenSource,
+            CancellationToken = cancellationTokenSource.Token,
             ChatGenerationRequest = request,
             Snapshot = snapshot,
             LastCheckpoint = DateTimeOffset.MinValue,
@@ -131,8 +137,23 @@ public sealed class AIJobService(
         if (!this.jobs.TryGetValue(jobId, out var job))
             return;
 
-        if (!job.CancellationTokenSource.IsCancellationRequested)
+        lock (job.SyncRoot)
-            await job.CancellationTokenSource.CancelAsync();
+        {
+            if (job.IsCompletionStarted)
+                return;
+        }
+
+        try
+        {
+            if (!job.CancellationTokenSource.IsCancellationRequested)
+                await job.CancellationTokenSource.CancelAsync();
+        }
+        catch (ObjectDisposedException)
+        {
+            return;
+        }
+
+        await this.CompleteChatGenerationAsync(job, AIJobStatus.CANCELED);
     }
 
     public async Task CancelChatGenerationAsync(Guid chatId)
@@ -167,13 +188,14 @@ public sealed class AIJobService(
     private async Task RunChatGenerationAsync(AIJobState state)
     {
         var request = state.ChatGenerationRequest;
-        var token = state.CancellationTokenSource.Token;
+        var token = state.CancellationToken;
 
         try
         {
+            token.ThrowIfCancellationRequested();
+
             var provider = request.ProviderSettings.CreateProvider();
             var chatThread = request.ChatThread;
-            var aiText = request.AIText;
 
             if (!chatThread.IsLLMProviderAllowed(provider))
             {
@@ -188,6 +210,8 @@ public sealed class AIJobService(
                 return;
             }
 
+            token.ThrowIfCancellationRequested();
+
             try
             {
                 var rag = new AISrcSelWithRetCtxVal();
@@ -207,21 +231,18 @@ public sealed class AIJobService(
                 logger.LogError(e, "Skipping the RAG process due to an error.");
             }
 
+            token.ThrowIfCancellationRequested();
+
             var lastStreamingEvent = DateTimeOffset.MinValue;
-            aiText.InitialRemoteWait = true;
+            if (!TrySetWaitingForRemote(state, token))
+                return;
             
             await this.NotifyChangedAsync(state);
             await foreach (var contentStreamChunk in provider.StreamChatCompletion(request.ProviderSettings.Model, chatThread, settingsManager, token))
             {
-                if (token.IsCancellationRequested)
+                if (!TryApplyStreamChunk(state, contentStreamChunk, token))
                     break;
 
-                aiText.InitialRemoteWait = false;
-                aiText.IsStreaming = true;
-                aiText.Text += contentStreamChunk;
-                aiText.Sources.MergeSources(contentStreamChunk.Sources);
-
-                UpdateStatus(state, AIJobStatus.RUNNING);
                 var now = DateTimeOffset.Now;
                 if (!settingsManager.ConfigurationData.App.IsSavingEnergy || now - lastStreamingEvent > STREAMING_EVENT_MIN_TIME)
                 {
@@ -255,11 +276,21 @@ public sealed class AIJobService(
 
     private async Task CompleteChatGenerationAsync(AIJobState state, AIJobStatus status, string errorMessage = "")
     {
+        lock (state.SyncRoot)
+        {
+            if (state.IsCompletionStarted)
+                return;
+
+            state.IsCompletionStarted = true;
+        }
+
         var aiText = state.ChatGenerationRequest.AIText;
         aiText.InitialRemoteWait = false;
         aiText.IsStreaming = false;
         aiText.Text = aiText.Text.RemoveThinkTags().Trim();
 
+        RemoveEmptyAIResponse(state);
+
         lock (state.SyncRoot)
         {
             state.Snapshot = state.Snapshot with
@@ -290,18 +321,41 @@ public sealed class AIJobService(
             state.ChatGenerationRequest.ChatThread.Blocks.Remove(aiBlock);
     }
 
-    private static void UpdateStatus(AIJobState state, AIJobStatus status)
+    private static bool TrySetWaitingForRemote(AIJobState state, CancellationToken token)
     {
         lock (state.SyncRoot)
         {
-            if (state.Snapshot.Status == status)
+            if (state.IsCompletionStarted || token.IsCancellationRequested)
-                return;
+                return false;
 
-            state.Snapshot = state.Snapshot with
+            state.ChatGenerationRequest.AIText.InitialRemoteWait = true;
+            return true;
+        }
+    }
+
+    private static bool TryApplyStreamChunk(AIJobState state, ContentStreamChunk contentStreamChunk, CancellationToken token)
+    {
+        lock (state.SyncRoot)
+        {
+            if (state.IsCompletionStarted || token.IsCancellationRequested)
+                return false;
+
+            var aiText = state.ChatGenerationRequest.AIText;
+            aiText.InitialRemoteWait = false;
+            aiText.IsStreaming = true;
+            aiText.Text += contentStreamChunk;
+            aiText.Sources.MergeSources(contentStreamChunk.Sources);
+
+            if (state.Snapshot.Status is not AIJobStatus.RUNNING)
             {
-                Status = status,
+                state.Snapshot = state.Snapshot with
-                UpdatedAt = DateTimeOffset.Now,
+                {
-            };
+                    Status = AIJobStatus.RUNNING,
+                    UpdatedAt = DateTimeOffset.Now,
+                };
+            }
+
+            return true;
         }
     }