Added support for parallel AI job processing (#774)

app/MindWork AI Studio/Components/ChatComponent.razor

@@ -68,7 +68,7 @@
             @if (this.SettingsManager.ConfigurationData.Workspace.StorageBehavior is WorkspaceStorageBehavior.STORE_CHATS_MANUALLY)
             {
                 <MudTooltip Text="@T("Save chat")" Placement="@TOOLBAR_TOOLTIP_PLACEMENT">
-                    <MudIconButton Icon="@Icons.Material.Filled.Save" OnClick="@(() => this.SaveThread())" Disabled="@(!this.CanThreadBeSaved || this.isStreaming)"/>
+                    <MudIconButton Icon="@Icons.Material.Filled.Save" OnClick="@(() => this.SaveThread())" Disabled="@(!this.CanThreadBeSaved || this.IsCurrentChatStreaming)"/>
                 </MudTooltip>
             }
 
@@ -89,14 +89,14 @@
             @if (this.SettingsManager.ConfigurationData.Workspace.StorageBehavior is WorkspaceStorageBehavior.STORE_CHATS_AUTOMATICALLY)
             {
                 <MudTooltip Text="@T("Delete this chat & start a new one.")" Placement="@TOOLBAR_TOOLTIP_PLACEMENT">
-                    <MudIconButton Icon="@Icons.Material.Filled.Refresh" OnClick="@(() => this.StartNewChat(useSameWorkspace: true, deletePreviousChat: true))" Disabled="@(!this.CanThreadBeSaved)"/>
+                    <MudIconButton Icon="@Icons.Material.Filled.Refresh" OnClick="@(() => this.StartNewChat(useSameWorkspace: true, deletePreviousChat: true))" Disabled="@(!this.CanThreadBeSaved || this.IsCurrentChatStreaming)"/>
                 </MudTooltip>
             }
 
             @if (this.SettingsManager.ConfigurationData.Workspace.StorageBehavior is not WorkspaceStorageBehavior.DISABLE_WORKSPACES)
             {
                 <MudTooltip Text="@T("Move the chat to a workspace, or to another if it is already in one.")" Placement="@TOOLBAR_TOOLTIP_PLACEMENT">
-                    <MudIconButton Icon="@Icons.Material.Filled.MoveToInbox" Disabled="@(!this.CanThreadBeSaved)" OnClick="@(() => this.MoveChatToWorkspace())"/>
+                    <MudIconButton Icon="@Icons.Material.Filled.MoveToInbox" Disabled="@(!this.CanThreadBeSaved || this.IsCurrentChatStreaming)" OnClick="@(() => this.MoveChatToWorkspace())"/>
                 </MudTooltip>
             }
             
@@ -134,7 +134,7 @@
                 <ConfidenceInfo Mode="PopoverTriggerMode.ICON" LLMProvider="@this.Provider.UsedLLMProvider"/>
             }
 
-            @if (this.isStreaming && this.cancellationTokenSource is not null)
+            @if (this.IsCurrentChatStreaming)
             {
                 <MudTooltip Text="@T("Stop generation")" Placement="@TOOLBAR_TOOLTIP_PLACEMENT">
                     <MudIconButton Icon="@Icons.Material.Filled.Stop" Color="Color.Error" OnClick="@(() => this.CancelStreaming())"/>

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

@@ -3,6 +3,7 @@ using AIStudio.Dialogs;
 using AIStudio.Provider;
 using AIStudio.Settings;
 using AIStudio.Settings.DataModel;
+using AIStudio.Tools.AIJobs;
 
 using Microsoft.AspNetCore.Components;
 using Microsoft.AspNetCore.Components.Web;
@@ -47,6 +48,9 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
     [Inject]
     private IJSRuntime JsRuntime { get; init; } = null!;
 
+    [Inject]
+    private AIJobService AIJobService { get; init; } = null!;
+
     private const Placement TOOLBAR_TOOLTIP_PLACEMENT = Placement.Top;
     private static readonly Dictionary<string, object?> USER_INPUT_ATTRIBUTES = new();
 
@@ -58,7 +62,6 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
     private bool mustScrollToBottomAfterRender;
     private InnerScrolling scrollingArea = null!;
     private byte scrollRenderCountdown;
-    private bool isStreaming;
     private string userInput = string.Empty;
     private bool mustStoreChat;
     private bool mustLoadChat;
@@ -67,8 +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 foregroundChatId = Guid.Empty;
     private int workspaceHeaderSyncVersion;
-    private CancellationTokenSource? cancellationTokenSource;
     private HashSet<FileAttachment> chatDocumentPaths = [];
 
     // Unfortunately, we need the input field reference to blur the focus away. Without
@@ -80,7 +83,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 ]);
+        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 ]);
         
         // Configure the spellchecking for the user input:
         this.SettingsManager.InjectSpellchecking(USER_INPUT_ATTRIBUTES);
@@ -217,6 +220,7 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
         
         // Select the correct provider:
         await this.SelectProviderWhenLoadingChat();
+        await this.SyncForegroundChatAsync();
         await base.OnInitializedAsync();
     }
 
@@ -273,6 +277,7 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
     protected override async Task OnParametersSetAsync()
     {
         await this.SyncWorkspaceHeaderWithChatThreadAsync();
+        await this.SyncForegroundChatAsync();
         await base.OnParametersSetAsync();
     }
 
@@ -333,7 +338,23 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
         this.WorkspaceName(this.currentWorkspaceName);
     }
 
+    private async Task SyncForegroundChatAsync()
+    {
+        var nextForegroundChatId = this.ChatThread?.ChatId ?? Guid.Empty;
+        if (this.foregroundChatId == nextForegroundChatId)
+            return;
+
+        if (this.foregroundChatId != Guid.Empty)
+            await this.AIJobService.SetForegroundAsync(AIJobKind.CHAT_GENERATION, this.foregroundChatId, false);
+
+        this.foregroundChatId = nextForegroundChatId;
+        if (this.foregroundChatId != Guid.Empty)
+            await this.AIJobService.SetForegroundAsync(AIJobKind.CHAT_GENERATION, this.foregroundChatId, true);
+    }
+
     private bool IsProviderSelected => this.Provider.UsedLLMProvider != LLMProviders.NONE;
+
+    private bool IsCurrentChatStreaming => this.ChatThread is not null && this.AIJobService.IsChatGenerationActive(this.ChatThread.ChatId);
     
     private string ProviderPlaceholder => this.IsProviderSelected ? T("Type your input here...") : T("Select a provider first");
 
@@ -453,7 +474,7 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
         if (!this.IsProviderSelected)
             return true;
         
-        if(this.isStreaming)
+        if(this.IsCurrentChatStreaming)
             return true;
         
         if(!this.ChatThread.IsLLMProviderAllowed(this.Provider))
@@ -614,7 +635,6 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
         await this.inputField.BlurAsync();
         
         // Enable the stream state for the chat component:
-        this.isStreaming = true;
         this.hasUnsavedChanges = true;
         
         if (this.SettingsManager.ConfigurationData.Chat.ShowLatestMessageAfterLoading)
@@ -624,38 +644,23 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
         }
         
         this.Logger.LogDebug($"Start processing user input using provider '{this.Provider.InstanceName}' with model '{this.Provider.Model}'.");
-        
-        using (this.cancellationTokenSource = new())
+        await this.AIJobService.TryStartChatGenerationAsync(new ChatGenerationRequest
         {
-            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.Provider.CreateProvider(), this.Provider.Model, lastUserPrompt, this.ChatThread, this.cancellationTokenSource.Token);
-        }
-        
-        this.cancellationTokenSource = null;
+            ChatThread = this.ChatThread!,
+            AIText = aiText,
+            LastUserPrompt = lastUserPrompt,
+            ProviderSettings = this.Provider,
+            IsForeground = true,
+        });
 
-        // Save the chat:
-        if (this.SettingsManager.ConfigurationData.Workspace.StorageBehavior is WorkspaceStorageBehavior.STORE_CHATS_AUTOMATICALLY)
-        {
-            await this.SaveThread();
-            this.hasUnsavedChanges = false;
-        }
-
-        // Disable the stream state:
-        this.isStreaming = false;
-        
-        // Update the UI:
+        await this.SyncForegroundChatAsync();
         this.StateHasChanged();
     }
     
     private async Task CancelStreaming()
     {
-        if (this.cancellationTokenSource is not null)
-            if(!this.cancellationTokenSource.IsCancellationRequested)
-                await this.cancellationTokenSource.CancelAsync();
+        if (this.ChatThread is not null)
+            await this.AIJobService.CancelChatGenerationAsync(this.ChatThread.ChatId);
     }
     
     private async Task SaveThread()
@@ -685,7 +690,7 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
         // Want the user to manage the chat storage manually? In that case, we have to ask the user
         // about possible data loss:
         //
-        if (this.SettingsManager.ConfigurationData.Workspace.StorageBehavior is WorkspaceStorageBehavior.STORE_CHATS_MANUALLY && this.hasUnsavedChanges)
+        if (this.SettingsManager.ConfigurationData.Workspace.StorageBehavior is WorkspaceStorageBehavior.STORE_CHATS_MANUALLY && this.hasUnsavedChanges && !this.IsCurrentChatStreaming)
         {
             var dialogParameters = new DialogParameters<ConfirmDialog>
             {
@@ -718,7 +723,6 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
         //
         // Reset our state:
         //
-        this.isStreaming = false;
         this.hasUnsavedChanges = false;
         this.userInput = string.Empty;
         
@@ -788,6 +792,7 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
         this.ApplyStandardDataSourceOptions();
         
         // Notify the parent component about the change:
+        await this.SyncForegroundChatAsync();
         await this.ChatThreadChanged.InvokeAsync(this.ChatThread);
     }
     
@@ -796,7 +801,7 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
         if(this.ChatThread is null)
             return;
         
-        if (this.SettingsManager.ConfigurationData.Workspace.StorageBehavior is WorkspaceStorageBehavior.STORE_CHATS_MANUALLY && this.hasUnsavedChanges)
+        if (this.SettingsManager.ConfigurationData.Workspace.StorageBehavior is WorkspaceStorageBehavior.STORE_CHATS_MANUALLY && this.hasUnsavedChanges && !this.IsCurrentChatStreaming)
         {
             var confirmationDialogParameters = new DialogParameters<ConfirmDialog>
             {
@@ -836,18 +841,21 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
     
     private async Task LoadedChatChanged()
     {
-        this.isStreaming = false;
         this.hasUnsavedChanges = false;
         this.userInput = string.Empty;
 
         if (this.ChatThread is not null)
         {
+            this.ChatThread = this.AIJobService.TryGetLiveChatThread(this.ChatThread.ChatId) ?? this.ChatThread;
+            await this.ChatThreadChanged.InvokeAsync(this.ChatThread);
             await this.SyncWorkspaceHeaderWithChatThreadAsync();
+            await this.SyncForegroundChatAsync();
             this.dataSourceSelectionComponent?.ChangeOptionWithoutSaving(this.ChatThread.DataSourceOptions, this.ChatThread.AISelectedDataSources);
         }
         else
         {
             this.ClearWorkspaceHeaderState();
+            await this.SyncForegroundChatAsync();
             this.ApplyStandardDataSourceOptions();
         }
         
@@ -863,12 +871,12 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
     
     private async Task ResetState()
     {
-        this.isStreaming = false;
         this.hasUnsavedChanges = false;
         this.userInput = string.Empty;
         this.ClearWorkspaceHeaderState();
         
         this.ChatThread = null;
+        await this.SyncForegroundChatAsync();
         this.ApplyStandardDataSourceOptions();
         await this.ChatThreadChanged.InvokeAsync(this.ChatThread);
     }
@@ -995,6 +1003,19 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
             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:
+                if (data is AIJobSnapshot { Kind: AIJobKind.CHAT_GENERATION } snapshot && this.ChatThread?.ChatId == snapshot.SubjectId)
+                {
+                    this.ChatThread = this.AIJobService.TryGetLiveChatThread(snapshot.SubjectId) ?? this.ChatThread;
+                    if (!snapshot.IsActive)
+                        this.hasUnsavedChanges = false;
+
+                    this.StateHasChanged();
+                }
+                break;
         }
     }
 
@@ -1005,6 +1026,9 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
             case Event.HAS_CHAT_UNSAVED_CHANGES:
                 if(this.SettingsManager.ConfigurationData.Workspace.StorageBehavior is WorkspaceStorageBehavior.STORE_CHATS_AUTOMATICALLY)
                     return Task.FromResult((TResult?) (object) false);
+
+                if (this.IsCurrentChatStreaming)
+                    return Task.FromResult((TResult?) (object) false);
                 
                 return Task.FromResult((TResult?)(object)this.hasUnsavedChanges);
         }
@@ -1024,21 +1048,8 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
             this.hasUnsavedChanges = false;
         }
 
-        if (this.cancellationTokenSource is not null)
-        {
-            try
-            {
-                if(!this.cancellationTokenSource.IsCancellationRequested)
-                    await this.cancellationTokenSource.CancelAsync();
-            
-                this.cancellationTokenSource.Dispose();
-            }
-            catch
-            {
-                // ignored
-            }
-        }
+        await this.AIJobService.SetForegroundAsync(AIJobKind.CHAT_GENERATION, this.foregroundChatId, false);
     }
 
     #endregion
-}
+}

app/MindWork AI Studio/Components/TreeItemData.cs

@@ -12,10 +12,16 @@ public class TreeItemData : ITreeItem
 
     public string Icon { get; init; } = string.Empty;
 
+    public string DefaultIcon { get; init; } = string.Empty;
+
     public TreeItemType Type { get; init; }
 
     public string Path { get; init; } = string.Empty;
 
+    public Guid ChatId { get; init; }
+
+    public Guid WorkspaceId { get; init; }
+
     public bool Expandable { get; init; } = true;
 
     public DateTimeOffset LastEditTime { get; init; }

app/MindWork AI Studio/Components/Workspaces.razor

@@ -24,7 +24,7 @@ else
                 case TreeItemData treeItem:
                     @if (treeItem.Type is TreeItemType.LOADING)
                     {
-                        <MudTreeViewItem T="ITreeItem" Icon="@treeItem.Icon" Value="@item.Value" Expanded="@item.Expanded" CanExpand="@false" Items="@(treeItem.Children!)">
+                        <MudTreeViewItem T="ITreeItem" Icon="@this.GetTreeItemIcon(treeItem)" Value="@item.Value" Expanded="@item.Expanded" CanExpand="@false" Items="@(treeItem.Children!)">
                             <BodyContent>
                                 <MudSkeleton Width="85%" Height="22px"/>
                             </BodyContent>
@@ -32,7 +32,7 @@ else
                     }
                     else if (treeItem.Type is TreeItemType.CHAT)
                     {
-                        <MudTreeViewItem T="ITreeItem" Icon="@treeItem.Icon" Value="@item.Value" Expanded="@item.Expanded" CanExpand="@treeItem.Expandable" Items="@(treeItem.Children!)" OnClick="@(() => this.LoadChatAsync(treeItem.Path, true))">
+                        <MudTreeViewItem T="ITreeItem" Icon="@this.GetTreeItemIcon(treeItem)" Value="@item.Value" Expanded="@item.Expanded" CanExpand="@treeItem.Expandable" Items="@(treeItem.Children!)" OnClick="@(() => this.LoadChatAsync(treeItem.Path, true))">
                             <BodyContent>
                                 <div style="display: grid; grid-template-columns: 1fr auto; align-items: center; width: 100%">
                                     <MudText Style="justify-self: start;">
@@ -48,15 +48,15 @@ else
                                     <div style="justify-self: end;">
 
                                         <MudTooltip Text="@T("Move to workspace")" Placement="@WORKSPACE_ITEM_TOOLTIP_PLACEMENT">
-                                            <MudIconButton Icon="@Icons.Material.Filled.MoveToInbox" Size="Size.Medium" Color="Color.Inherit" OnClick="@(() => this.MoveChatAsync(treeItem.Path))"/>
+                                            <MudIconButton Icon="@Icons.Material.Filled.MoveToInbox" Size="Size.Medium" Color="Color.Inherit" Disabled="@this.IsChatTreeItemBusy(treeItem)" OnClick="@(() => this.MoveChatAsync(treeItem.Path))"/>
                                         </MudTooltip>
 
                                         <MudTooltip Text="@T("Rename")" Placement="@WORKSPACE_ITEM_TOOLTIP_PLACEMENT">
-                                            <MudIconButton Icon="@Icons.Material.Filled.Edit" Size="Size.Medium" Color="Color.Inherit" OnClick="@(() => this.RenameChatAsync(treeItem.Path))"/>
+                                            <MudIconButton Icon="@Icons.Material.Filled.Edit" Size="Size.Medium" Color="Color.Inherit" Disabled="@this.IsChatTreeItemBusy(treeItem)" OnClick="@(() => this.RenameChatAsync(treeItem.Path))"/>
                                         </MudTooltip>
 
                                         <MudTooltip Text="@T("Delete")" Placement="@WORKSPACE_ITEM_TOOLTIP_PLACEMENT">
-                                            <MudIconButton Icon="@Icons.Material.Filled.Delete" Size="Size.Medium" Color="Color.Error" OnClick="@(() => this.DeleteChatAsync(treeItem.Path))"/>
+                                            <MudIconButton Icon="@Icons.Material.Filled.Delete" Size="Size.Medium" Color="Color.Error" Disabled="@this.IsChatTreeItemBusy(treeItem)" OnClick="@(() => this.DeleteChatAsync(treeItem.Path))"/>
                                         </MudTooltip>
                                     </div>
                                 </div>
@@ -65,7 +65,7 @@ else
                     }
                     else if (treeItem.Type is TreeItemType.WORKSPACE)
                     {
-                        <MudTreeViewItem T="ITreeItem" Icon="@treeItem.Icon" Value="@item.Value" Expanded="@item.Expanded" CanExpand="@treeItem.Expandable" Items="@(treeItem.Children!)" OnClick="@(() => this.OnWorkspaceClicked(treeItem))">
+                        <MudTreeViewItem T="ITreeItem" Icon="@this.GetTreeItemIcon(treeItem)" Value="@item.Value" Expanded="@item.Expanded" CanExpand="@treeItem.Expandable" Items="@(treeItem.Children!)" OnClick="@(() => this.OnWorkspaceClicked(treeItem))">
                             <BodyContent>
                                 <div style="display: grid; grid-template-columns: 1fr auto; align-items: center; width: 100%">
                                     <MudText Style="justify-self: start;">
@@ -86,7 +86,7 @@ else
                     }
                     else
                     {
-                        <MudTreeViewItem T="ITreeItem" Icon="@treeItem.Icon" Value="@item.Value" Expanded="@item.Expanded" CanExpand="@treeItem.Expandable" Items="@(treeItem.Children!)">
+                        <MudTreeViewItem T="ITreeItem" Icon="@this.GetTreeItemIcon(treeItem)" Value="@item.Value" Expanded="@item.Expanded" CanExpand="@treeItem.Expandable" Items="@(treeItem.Children!)">
                             <BodyContent>
                                 <div style="display: grid; grid-template-columns: 1fr auto; align-items: center; width: 100%">
                                     <MudText Style="justify-self: start;">

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

@@ -4,6 +4,7 @@ using System.Text.Json;
 using AIStudio.Chat;
 using AIStudio.Dialogs;
 using AIStudio.Settings;
+using AIStudio.Tools.AIJobs;
 
 using Microsoft.AspNetCore.Components;
 
@@ -18,6 +19,9 @@ public partial class Workspaces : MSGComponentBase
     
     [Inject]
     private ILogger<Workspaces> Logger { get; init; } = null!;
+
+    [Inject]
+    private AIJobService AIJobService { get; init; } = null!;
     
     [Parameter]
     public ChatThread? CurrentChatThread { get; set; }
@@ -42,6 +46,7 @@ public partial class Workspaces : MSGComponentBase
     protected override async Task OnInitializedAsync()
     {
         await base.OnInitializedAsync();
+        this.ApplyFilters([], [ Event.AI_JOB_CHANGED, Event.AI_JOB_FINISHED, Event.CHAT_GENERATION_CHANGED ]);
         _ = this.LoadTreeItemsAsync(startPrefetch: true);
     }
 
@@ -111,7 +116,7 @@ public partial class Workspaces : MSGComponentBase
         
         var temporaryChatsChildren = new List<TreeItemData<ITreeItem>>();
         foreach (var temporaryChat in snapshot.TemporaryChats.OrderByDescending(x => x.LastEditTime))
-            temporaryChatsChildren.Add(CreateChatTreeItem(temporaryChat, WorkspaceBranch.TEMPORARY_CHATS, depth: 1, icon: Icons.Material.Filled.Timer));
+            temporaryChatsChildren.Add(this.CreateChatTreeItem(temporaryChat, WorkspaceBranch.TEMPORARY_CHATS, depth: 1, icon: Icons.Material.Filled.Timer));
 
         this.treeItems.Add(new TreeItemData<ITreeItem>
         {
@@ -136,7 +141,7 @@ public partial class Workspaces : MSGComponentBase
         if (workspace.ChatsLoaded)
         {
             foreach (var workspaceChat in workspace.Chats.OrderByDescending(x => x.LastEditTime))
-                children.Add(CreateChatTreeItem(workspaceChat, WorkspaceBranch.WORKSPACES, depth: 2, icon: Icons.Material.Filled.Chat));
+                children.Add(this.CreateChatTreeItem(workspaceChat, WorkspaceBranch.WORKSPACES, depth: 2, icon: Icons.Material.Filled.Chat));
         }
         else if (this.loadingWorkspaceChatLists.Contains(workspace.WorkspaceId))
             children.AddRange(this.CreateLoadingRows(workspace.WorkspacePath));
@@ -192,7 +197,7 @@ public partial class Workspaces : MSGComponentBase
         };
     }
 
-    private static TreeItemData<ITreeItem> CreateChatTreeItem(WorkspaceTreeChat chat, WorkspaceBranch branch, int depth, string icon)
+    private TreeItemData<ITreeItem> CreateChatTreeItem(WorkspaceTreeChat chat, WorkspaceBranch branch, int depth, string icon)
     {
         return new TreeItemData<ITreeItem>
         {
@@ -204,13 +209,43 @@ public partial class Workspaces : MSGComponentBase
                 Branch = branch,
                 Text = chat.Name,
                 Icon = icon,
+                DefaultIcon = icon,
                 Expandable = false,
                 Path = chat.ChatPath,
+                ChatId = chat.ChatId,
+                WorkspaceId = chat.WorkspaceId,
                 LastEditTime = chat.LastEditTime,
             },
         };
     }
 
+    private string GetTreeItemIcon(TreeItemData treeItem)
+    {
+        if (treeItem.Type is not TreeItemType.CHAT)
+            return treeItem.Icon;
+
+        var defaultIcon = string.IsNullOrWhiteSpace(treeItem.DefaultIcon) ? treeItem.Icon : treeItem.DefaultIcon;
+        return this.GetChatTreeIcon(treeItem.ChatId, defaultIcon);
+    }
+
+    private bool IsChatTreeItemBusy(TreeItemData treeItem)
+    {
+        return treeItem.Type is TreeItemType.CHAT && this.AIJobService.IsChatGenerationActive(treeItem.ChatId);
+    }
+
+    private string GetChatTreeIcon(Guid chatId, string defaultIcon)
+    {
+        var snapshot = this.AIJobService.TryGetChatSnapshot(chatId);
+        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,
+        };
+    }
+
     private async Task SafeStateHasChanged()
     {
         if (this.isDisposed)
@@ -348,6 +383,9 @@ public partial class Workspaces : MSGComponentBase
         {
             var chatData = await File.ReadAllTextAsync(Path.Join(chatPath, "thread.json"), Encoding.UTF8);
             var chat = JsonSerializer.Deserialize<ChatThread>(chatData, WorkspaceBehaviour.JSON_OPTIONS);
+            if (chat is not null)
+                chat = this.AIJobService.TryGetLiveChatThread(chat.ChatId) ?? chat;
+
             if (switchToChat)
             {
                 this.CurrentChatThread = chat;
@@ -371,6 +409,9 @@ public partial class Workspaces : MSGComponentBase
         if (chat is null)
             return;
 
+        if (this.AIJobService.IsChatGenerationActive(chat.ChatId))
+            return;
+
         if (askForConfirmation)
         {
             var workspaceName = await WorkspaceBehaviour.LoadWorkspaceNameAsync(chat.WorkspaceId);
@@ -407,6 +448,9 @@ public partial class Workspaces : MSGComponentBase
         var chat = await this.LoadChatAsync(chatPath, false);
         if (chat is null)
             return;
+
+        if (this.AIJobService.IsChatGenerationActive(chat.ChatId))
+            return;
         
         var dialogParameters = new DialogParameters<SingleInputDialog>
         {
@@ -525,6 +569,9 @@ public partial class Workspaces : MSGComponentBase
         var chat = await this.LoadChatAsync(chatPath, false);
         if (chat is null)
             return;
+
+        if (this.AIJobService.IsChatGenerationActive(chat.ChatId))
+            return;
         
         var dialogParameters = new DialogParameters<WorkspaceSelectionDialog>
         {
@@ -597,6 +644,12 @@ public partial class Workspaces : MSGComponentBase
             case Event.PLUGINS_RELOADED:
                 await this.ForceRefreshFromDiskAsync();
                 break;
+
+            case Event.AI_JOB_CHANGED:
+            case Event.AI_JOB_FINISHED:
+            case Event.CHAT_GENERATION_CHANGED:
+                await this.SafeStateHasChanged();
+                break;
         }
     }