Fix diff views with background assistants

app/MindWork AI Studio/Assistants/AssistantBase.razor.cs

@@ -142,6 +142,7 @@ public abstract partial class AssistantBase<TSettings> : AssistantLowerBase wher
     private bool isDisposed;
     private AssistantSessionKey assistantSessionKey;
     private Guid? assistantSessionId;
+    private AssistantSessionSnapshot? pendingRenderedAssistantSessionSnapshot;
 
     /// <summary>
     /// Gets whether the Blazor component instance has already been disposed.
@@ -195,6 +196,12 @@ public abstract partial class AssistantBase<TSettings> : AssistantLowerBase wher
         // We don't want to show validation errors when the user opens the dialog.
         if(firstRender)
             this.Form?.ResetValidation();
+
+        if (this.pendingRenderedAssistantSessionSnapshot is { } snapshot)
+        {
+            this.pendingRenderedAssistantSessionSnapshot = null;
+            await this.OnAssistantSessionRenderedAsync(snapshot);
+        }
         
         await base.OnAfterRenderAsync(firstRender);
     }
@@ -229,7 +236,7 @@ public abstract partial class AssistantBase<TSettings> : AssistantLowerBase wher
 
         this.CancellationTokenSource = new();
         this.isProcessing = true;
-        var startedSession = this.AssistantSessionService.TryBegin(this.assistantSessionKey, this.Title, this.CancellationTokenSource, this.ChatThread, this.CaptureAssistantSessionState());
+        var startedSession = await this.AssistantSessionService.TryBeginAsync(this.assistantSessionKey, this.Title, this.CancellationTokenSource, this.ChatThread, this.CaptureAssistantSessionState());
         if (startedSession.IsActive is not true || startedSession.Key != this.assistantSessionKey)
         {
             this.CancellationTokenSource.Dispose();
@@ -658,6 +665,13 @@ public abstract partial class AssistantBase<TSettings> : AssistantLowerBase wher
     /// <returns>A task that completes after derived UI restore work has finished.</returns>
     protected virtual Task OnAssistantSessionAttachedAsync(AssistantSessionSnapshot snapshot) => Task.CompletedTask;
 
+    /// <summary>
+    /// Allows derived assistants to restore DOM-dependent client-only UI after an attached session was rendered.
+    /// </summary>
+    /// <param name="snapshot">The assistant session snapshot that was rendered.</param>
+    /// <returns>A task that completes after derived UI restore work has finished.</returns>
+    protected virtual Task OnAssistantSessionRenderedAsync(AssistantSessionSnapshot snapshot) => Task.CompletedTask;
+
     /// <summary>
     /// Handles assistant session change events for the current assistant instance.
     /// </summary>
@@ -710,6 +724,9 @@ public abstract partial class AssistantBase<TSettings> : AssistantLowerBase wher
         if (restoreClientOnlyContent)
             await this.OnAssistantSessionAttachedAsync(snapshot);
 
+        if (restoreClientOnlyContent)
+            this.pendingRenderedAssistantSessionSnapshot = snapshot;
+
         await this.RefreshAssistantUIAsync();
     }
 

app/MindWork AI Studio/Assistants/GrammarSpelling/AssistantGrammarSpelling.razor.cs

@@ -154,7 +154,7 @@ public partial class AssistantGrammarSpelling : AssistantBaseCore<SettingsDialog
             await this.JsRuntime.GenerateAndShowDiff(this.inputText, this.correctedText);
     }
 
-    protected override async Task OnAssistantSessionAttachedAsync(AssistantSessionSnapshot snapshot)
+    protected override async Task OnAssistantSessionRenderedAsync(AssistantSessionSnapshot snapshot)
     {
         if (!snapshot.IsActive && !string.IsNullOrWhiteSpace(this.inputText) && !string.IsNullOrWhiteSpace(this.correctedText))
             await this.JsRuntime.GenerateAndShowDiff(this.inputText, this.correctedText);

app/MindWork AI Studio/Assistants/RewriteImprove/AssistantRewriteImprove.razor.cs

@@ -167,7 +167,7 @@ public partial class AssistantRewriteImprove : AssistantBaseCore<SettingsDialogR
             await this.JsRuntime.GenerateAndShowDiff(this.inputText, this.rewrittenText);
     }
 
-    protected override async Task OnAssistantSessionAttachedAsync(AssistantSessionSnapshot snapshot)
+    protected override async Task OnAssistantSessionRenderedAsync(AssistantSessionSnapshot snapshot)
     {
         if (!snapshot.IsActive && !string.IsNullOrWhiteSpace(this.inputText) && !string.IsNullOrWhiteSpace(this.rewrittenText))
             await this.JsRuntime.GenerateAndShowDiff(this.inputText, this.rewrittenText);

app/MindWork AI Studio/Tools/AssistantSessions/AssistantSessionService.cs

@@ -122,7 +122,7 @@ public sealed class AssistantSessionService(MessageBus messageBus)
     /// <param name="chatThread">The current assistant chat thread, if one already exists.</param>
     /// <param name="state">The initial assistant component state.</param>
     /// <returns>The new session snapshot, or the existing active session snapshot.</returns>
-    public AssistantSessionSnapshot TryBegin(AssistantSessionKey key, string title, CancellationTokenSource cancellationTokenSource, ChatThread? chatThread, Dictionary<string, IAssistantSessionSnapshotField> state)
+    public async Task<AssistantSessionSnapshot> TryBeginAsync(AssistantSessionKey key, string title, CancellationTokenSource cancellationTokenSource, ChatThread? chatThread, Dictionary<string, IAssistantSessionSnapshotField> state)
     {
         if (this.sessions.TryGetValue(key, out var existing) && existing.Status is AssistantSessionStatus.RUNNING or AssistantSessionStatus.CANCELING)
             return CreateSnapshot(existing);
@@ -142,8 +142,9 @@ public sealed class AssistantSessionService(MessageBus messageBus)
         };
 
         this.sessions[key] = session;
-        _ = this.NotifyChangedAsync(session);
+        var snapshot = CreateSnapshot(session);
-        return CreateSnapshot(session);
+        await this.NotifyChangedAsync(session);
+        return snapshot;
     }
 
     /// <summary>

app/MindWork AI Studio/wwwroot/app.js

@@ -1,10 +1,18 @@
 window.generateDiff = function (text1, text2, divDiff, divLegend) {
     let wikEdDiff = new WikEdDiff();
     let targetDiv = document.getElementById(divDiff)
+    if (!targetDiv) {
+        return;
+    }
+
     targetDiv.innerHTML = wikEdDiff.diff(text1, text2);
     targetDiv.classList.add('mud-typography-body1', 'improvedDiff');
     
     let legend = document.getElementById(divLegend);
+    if (!legend) {
+        return;
+    }
+
     legend.innerHTML = `
     <div class="legend mt-2">
         <h3>Legend</h3>
@@ -20,6 +28,10 @@ window.generateDiff = function (text1, text2, divDiff, divLegend) {
 
 window.clearDiv = function (divName) {
     let targetDiv = document.getElementById(divName);
+    if (!targetDiv) {
+        return;
+    }
+
     targetDiv.innerHTML = '';
 }