Optimized rendering performance

app/MindWork AI Studio/App.razor

@@ -27,6 +27,7 @@
         <script src="system/MudBlazor.Markdown/MudBlazor.Markdown.min.js"></script>
         <script src="system/CodeBeam.MudBlazor.Extensions/MudExtensions.min.js"></script>
         <script src="app.js"></script>
+        <script src="chat-math.js"></script>
         <script src="audio.js"></script>
     </body>
 

app/MindWork AI Studio/Chat/ContentBlockComponent.razor

@@ -97,22 +97,24 @@
                             else
                             {
                                 var renderPlan = this.GetMarkdownRenderPlan(textContent.Text);
-                                foreach (var segment in renderPlan.Segments)
-                                {
-                                    var segmentContent = segment.GetContent(renderPlan.Source);
-                                    if (segment.Type is MarkdownRenderSegmentType.MARKDOWN)
+                                <div @ref="this.mathContentContainer" class="chat-math-container">
+                                    @foreach (var segment in renderPlan.Segments)
                                     {
-                                        <MudMarkdown Value="@segmentContent" Props="Markdown.DefaultConfig" Styling="@this.MarkdownStyling" MarkdownPipeline="Markdown.SAFE_MARKDOWN_PIPELINE" />
+                                        var segmentContent = segment.GetContent(renderPlan.Source);
+                                        if (segment.Type is MarkdownRenderSegmentType.MARKDOWN)
+                                        {
+                                            <MudMarkdown @key="@segment.RenderKey" Value="@segmentContent" Props="Markdown.DefaultConfig" Styling="@this.MarkdownStyling" MarkdownPipeline="Markdown.SAFE_MARKDOWN_PIPELINE" />
+                                        }
+                                        else
+                                        {
+                                            <MathJaxBlock @key="@segment.RenderKey" Value="@segmentContent" Class="mb-5" />
+                                        }
                                     }
-                                    else
+                                    @if (textContent.Sources.Count > 0)
                                     {
-                                        <MathJaxBlock Value="@segmentContent" Class="mb-5" />
+                                        <MudMarkdown Value="@textContent.Sources.ToMarkdown()" Props="Markdown.DefaultConfig" Styling="@this.MarkdownStyling" MarkdownPipeline="Markdown.SAFE_MARKDOWN_PIPELINE" />
                                     }
-                                }
-                                @if (textContent.Sources.Count > 0)
-                                {
-                                    <MudMarkdown Value="@textContent.Sources.ToMarkdown()" Props="Markdown.DefaultConfig" Styling="@this.MarkdownStyling" MarkdownPipeline="Markdown.SAFE_MARKDOWN_PIPELINE" />
-                                }
+                                </div>
                             }
                         }
                     }
@@ -147,4 +149,4 @@
             }
         }
     </MudCardContent>
-</MudCard>
+</MudCard>

app/MindWork AI Studio/Chat/ContentBlockComponent.razor.cs

@@ -8,8 +8,10 @@ namespace AIStudio.Chat;
 /// <summary>
 /// The UI component for a chat content block, i.e., for any IContent.
 /// </summary>
-public partial class ContentBlockComponent : MSGComponentBase
+public partial class ContentBlockComponent : MSGComponentBase, IAsyncDisposable
 {
+    private const string CHAT_MATH_SYNC_FUNCTION = "chatMath.syncContainer";
+    private const string CHAT_MATH_DISPOSE_FUNCTION = "chatMath.disposeContainer";
     private const string HTML_START_TAG = "<";
     private const string HTML_END_TAG = "</";
     private const string HTML_SELF_CLOSING_TAG = "/>";
@@ -89,11 +91,18 @@ public partial class ContentBlockComponent : MSGComponentBase
     [Inject]
     private RustService RustService { get; init; } = null!;
 
+    [Inject]
+    private IJSRuntime JsRuntime { get; init; } = null!;
+
     private bool HideContent { get; set; }
     private bool hasRenderHash;
     private int lastRenderHash;
     private string cachedMarkdownRenderPlanInput = string.Empty;
     private MarkdownRenderPlan cachedMarkdownRenderPlan = MarkdownRenderPlan.EMPTY;
+    private ElementReference mathContentContainer;
+    private string lastMathRenderSignature = string.Empty;
+    private bool hasActiveMathContainer;
+    private bool isDisposed;
 
     #region Overrides of ComponentBase
 
@@ -109,6 +118,12 @@ public partial class ContentBlockComponent : MSGComponentBase
         return base.OnParametersSetAsync();
     }
 
+    protected override async Task OnAfterRenderAsync(bool firstRender)
+    {
+        await this.SyncMathRenderIfNeededAsync();
+        await base.OnAfterRenderAsync(firstRender);
+    }
+
     /// <inheritdoc />
     protected override bool ShouldRender()
     {
@@ -216,6 +231,81 @@ public partial class ContentBlockComponent : MSGComponentBase
         return this.cachedMarkdownRenderPlan;
     }
 
+    private async Task SyncMathRenderIfNeededAsync()
+    {
+        if (this.isDisposed)
+            return;
+
+        if (!this.TryGetCompletedMathRenderState(out var mathRenderSignature))
+        {
+            await this.DisposeMathContainerIfNeededAsync();
+            return;
+        }
+
+        if (string.Equals(this.lastMathRenderSignature, mathRenderSignature, StringComparison.Ordinal))
+            return;
+
+        await this.JsRuntime.InvokeVoidAsync(CHAT_MATH_SYNC_FUNCTION, this.mathContentContainer, mathRenderSignature);
+        this.lastMathRenderSignature = mathRenderSignature;
+        this.hasActiveMathContainer = true;
+    }
+
+    private async Task DisposeMathContainerIfNeededAsync()
+    {
+        if (!this.hasActiveMathContainer)
+        {
+            this.lastMathRenderSignature = string.Empty;
+            return;
+        }
+
+        try
+        {
+            await this.JsRuntime.InvokeVoidAsync(CHAT_MATH_DISPOSE_FUNCTION, this.mathContentContainer);
+        }
+        catch (JSDisconnectedException)
+        {
+        }
+        catch (ObjectDisposedException)
+        {
+        }
+
+        this.hasActiveMathContainer = false;
+        this.lastMathRenderSignature = string.Empty;
+    }
+
+    private bool TryGetCompletedMathRenderState(out string mathRenderSignature)
+    {
+        mathRenderSignature = string.Empty;
+
+        if (this.HideContent || this.Type is not ContentType.TEXT || this.Content.IsStreaming || this.Content is not ContentText textContent || textContent.InitialRemoteWait)
+            return false;
+
+        var renderPlan = this.GetMarkdownRenderPlan(textContent.Text);
+        mathRenderSignature = CreateMathRenderSignature(renderPlan);
+        return !string.IsNullOrEmpty(mathRenderSignature);
+    }
+
+    private static string CreateMathRenderSignature(MarkdownRenderPlan renderPlan)
+    {
+        var hash = new HashCode();
+        var mathSegmentCount = 0;
+
+        foreach (var segment in renderPlan.Segments)
+        {
+            if (segment.Type is not MarkdownRenderSegmentType.MATH_BLOCK)
+                continue;
+
+            mathSegmentCount++;
+            hash.Add(segment.Start);
+            hash.Add(segment.Length);
+            hash.Add(segment.GetContent(renderPlan.Source).GetHashCode(StringComparison.Ordinal));
+        }
+
+        return mathSegmentCount == 0
+            ? string.Empty
+            : $"{mathSegmentCount}:{hash.ToHashCode()}";
+    }
+
     private static MarkdownRenderPlan BuildMarkdownRenderPlan(string text)
     {
         var normalized = NormalizeMarkdownForRendering(text);
@@ -428,9 +518,11 @@ public partial class ContentBlockComponent : MSGComponentBase
 
         public MarkdownRenderSegmentType Type { get; } = type;
 
-        private int Start { get; } = start;
+        public int Start { get; } = start;
 
-        private int Length { get; } = length;
+        public int Length { get; } = length;
+
+        public int RenderKey { get; } = HashCode.Combine(type, start, length);
 
         public string GetContent(string source)
         {
@@ -517,4 +609,14 @@ public partial class ContentBlockComponent : MSGComponentBase
         var result = await ReviewAttachmentsDialog.OpenDialogAsync(this.DialogService, this.Content.FileAttachments.ToHashSet());
         this.Content.FileAttachments = result.ToList();
     }
+
+    public async ValueTask DisposeAsync()
+    {
+        if (this.isDisposed)
+            return;
+
+        this.isDisposed = true;
+        await this.DisposeMathContainerIfNeededAsync();
+        this.Dispose();
+    }
 }

app/MindWork AI Studio/Chat/MathJaxBlock.razor

@@ -1,5 +1,5 @@
 @namespace AIStudio.Chat
 
-<div class="@this.RootClass" style="white-space: pre-wrap;">
+<div class="@this.RootClass" data-chat-math-block="true" style="white-space: pre-wrap;">
     @this.MathText
 </div>

app/MindWork AI Studio/Chat/MathJaxBlock.razor.cs

@@ -4,27 +4,15 @@ namespace AIStudio.Chat;
 
 public partial class MathJaxBlock
 {
-    private const string MATH_JAX_SCRIPT_ID = "mudblazor-markdown-mathjax";
-
     [Parameter]
     public string Value { get; init; } = string.Empty;
 
     [Parameter]
     public string Class { get; init; } = string.Empty;
 
-    [Inject]
-    private IJSRuntime JsRuntime { get; init; } = null!;
-
     private string RootClass => string.IsNullOrWhiteSpace(this.Class)
         ? "chat-mathjax-block"
         : $"chat-mathjax-block {this.Class}";
 
     private string MathText => $"$${Environment.NewLine}{this.Value}{Environment.NewLine}$$";
-
-    protected override async Task OnAfterRenderAsync(bool firstRender)
-    {
-        await this.JsRuntime.InvokeVoidAsync("appendMathJaxScript", MATH_JAX_SCRIPT_ID);
-        await this.JsRuntime.InvokeVoidAsync("refreshMathJaxScript");
-        await base.OnAfterRenderAsync(firstRender);
-    }
 }

app/MindWork AI Studio/wwwroot/chat-math.js

@@ -0,0 +1,280 @@
+const MATH_JAX_SCRIPT_ID = 'mudblazor-markdown-mathjax'
+const MATH_JAX_SCRIPT_SRC = '_content/MudBlazor.Markdown/MudBlazor.Markdown.MathJax.min.js'
+const INTERSECTION_ROOT_MARGIN = '240px 0px 240px 0px'
+const MAX_TYPES_PER_BATCH = 4
+const containerStates = new Map()
+const pendingMathElements = new Set()
+
+let mathJaxReadyPromise = null
+let batchScheduled = false
+let typesetInProgress = false
+
+function isMathJaxReady() {
+    return typeof window.MathJax?.typesetPromise === 'function' || typeof window.MathJax?.typeset === 'function'
+}
+
+function waitForMathJaxReady(attempt = 0) {
+    if (isMathJaxReady())
+        return Promise.resolve()
+
+    if (attempt >= 80)
+        return Promise.reject(new Error('MathJax did not finish loading in time.'))
+
+    return new Promise((resolve, reject) => {
+        window.setTimeout(() => {
+            waitForMathJaxReady(attempt + 1).then(resolve).catch(reject)
+        }, 50)
+    })
+}
+
+function ensureMathJaxLoaded() {
+    if (isMathJaxReady())
+        return Promise.resolve()
+
+    if (mathJaxReadyPromise)
+        return mathJaxReadyPromise
+
+    mathJaxReadyPromise = new Promise((resolve, reject) => {
+        let script = document.getElementById(MATH_JAX_SCRIPT_ID)
+
+        const onLoad = () => {
+            waitForMathJaxReady().then(resolve).catch(reject)
+        }
+
+        const onError = () => reject(new Error('Failed to load the MathJax script.'))
+
+        if (!script) {
+            script = document.createElement('script')
+            script.id = MATH_JAX_SCRIPT_ID
+            script.type = 'text/javascript'
+            script.src = MATH_JAX_SCRIPT_SRC
+            script.addEventListener('load', onLoad, { once: true })
+            script.addEventListener('error', onError, { once: true })
+            document.head.appendChild(script)
+            return
+        }
+
+        script.addEventListener('load', onLoad, { once: true })
+        script.addEventListener('error', onError, { once: true })
+        void waitForMathJaxReady().then(resolve).catch(() => {})
+    }).catch(error => {
+        mathJaxReadyPromise = null
+        throw error
+    })
+
+    return mathJaxReadyPromise
+}
+
+function createContainerState() {
+    return {
+        signature: '',
+        observer: null,
+        observedElements: new Set()
+    }
+}
+
+function disconnectContainerState(state) {
+    if (state.observer) {
+        state.observer.disconnect()
+        state.observer = null
+    }
+
+    for (const element of state.observedElements)
+        pendingMathElements.delete(element)
+
+    state.observedElements.clear()
+}
+
+function isNearViewport(element) {
+    const rect = element.getBoundingClientRect()
+    return rect.bottom >= -240 && rect.top <= window.innerHeight + 240
+}
+
+function queueElementForTypeset(element, signature) {
+    if (!element || !element.isConnected)
+        return
+
+    if (element.dataset.chatMathProcessedSignature === signature)
+        return
+
+    element.dataset.chatMathTargetSignature = signature
+    element.dataset.chatMathPending = 'true'
+    pendingMathElements.add(element)
+    schedulePendingTypeset(false)
+}
+
+function schedulePendingTypeset(useIdleCallback) {
+    if (batchScheduled)
+        return
+
+    batchScheduled = true
+    const flush = () => {
+        batchScheduled = false
+        void flushPendingTypeset()
+    }
+
+    if (useIdleCallback && typeof window.requestIdleCallback === 'function') {
+        window.requestIdleCallback(flush, { timeout: 120 })
+        return
+    }
+
+    window.requestAnimationFrame(flush)
+}
+
+async function flushPendingTypeset() {
+    if (typesetInProgress || pendingMathElements.size === 0)
+        return
+
+    typesetInProgress = true
+    const elementsToTypeset = []
+
+    try {
+        await ensureMathJaxLoaded()
+
+        for (const element of pendingMathElements) {
+            if (elementsToTypeset.length >= MAX_TYPES_PER_BATCH)
+                break
+
+            if (!element.isConnected) {
+                pendingMathElements.delete(element)
+                continue
+            }
+
+            const targetSignature = element.dataset.chatMathTargetSignature ?? ''
+            if (element.dataset.chatMathProcessedSignature === targetSignature) {
+                pendingMathElements.delete(element)
+                element.dataset.chatMathPending = 'false'
+                continue
+            }
+
+            elementsToTypeset.push(element)
+        }
+
+        if (elementsToTypeset.length === 0)
+            return
+
+        for (const element of elementsToTypeset)
+            pendingMathElements.delete(element)
+
+        if (typeof window.MathJax?.typesetClear === 'function') {
+            try {
+                window.MathJax.typesetClear(elementsToTypeset)
+            } catch (error) {
+                console.warn('chatMath: failed to clear previous MathJax state.', error)
+            }
+        }
+
+        if (typeof window.MathJax?.typesetPromise === 'function')
+            await window.MathJax.typesetPromise(elementsToTypeset)
+        else if (typeof window.MathJax?.typeset === 'function')
+            window.MathJax.typeset(elementsToTypeset)
+
+        for (const element of elementsToTypeset) {
+            element.dataset.chatMathProcessedSignature = element.dataset.chatMathTargetSignature ?? ''
+            element.dataset.chatMathPending = 'false'
+        }
+    } catch (error) {
+        console.warn('chatMath: failed to typeset math content.', error)
+
+        for (const element of elementsToTypeset)
+            if (element.isConnected)
+                pendingMathElements.add(element)
+    } finally {
+        typesetInProgress = false
+
+        if (pendingMathElements.size > 0)
+            schedulePendingTypeset(true)
+    }
+}
+
+function createIntersectionObserver(state, signature) {
+    return new IntersectionObserver(entries => {
+        let queuedVisibleElement = false
+
+        for (const entry of entries) {
+            if (!entry.isIntersecting)
+                continue
+
+            const element = entry.target
+            state.observer?.unobserve(element)
+            state.observedElements.delete(element)
+            queueElementForTypeset(element, signature)
+            queuedVisibleElement = true
+        }
+
+        if (queuedVisibleElement)
+            schedulePendingTypeset(true)
+    }, {
+        root: null,
+        rootMargin: INTERSECTION_ROOT_MARGIN,
+        threshold: 0.01
+    })
+}
+
+function getMathElements(container) {
+    return Array.from(container.querySelectorAll('.chat-mathjax-block'))
+}
+
+window.chatMath = {
+    syncContainer: async function(container, signature) {
+        if (!container)
+            return
+
+        let state = containerStates.get(container)
+        if (!state) {
+            state = createContainerState()
+            containerStates.set(container, state)
+        }
+
+        if (state.signature === signature)
+            return
+
+        disconnectContainerState(state)
+        state.signature = signature
+
+        const mathElements = getMathElements(container)
+        if (mathElements.length === 0)
+            return
+
+        await ensureMathJaxLoaded()
+
+        state.observer = createIntersectionObserver(state, signature)
+
+        for (const element of mathElements) {
+            if (isNearViewport(element)) {
+                queueElementForTypeset(element, signature)
+                continue
+            }
+
+            element.dataset.chatMathTargetSignature = signature
+            state.observer.observe(element)
+            state.observedElements.add(element)
+        }
+
+        schedulePendingTypeset(false)
+    },
+
+    disposeContainer: function(container) {
+        if (!container)
+            return
+
+        const state = containerStates.get(container)
+        if (!state)
+            return
+
+        disconnectContainerState(state)
+        containerStates.delete(container)
+
+        const mathElements = getMathElements(container)
+        for (const element of mathElements)
+            pendingMathElements.delete(element)
+
+        if (typeof window.MathJax?.typesetClear === 'function' && mathElements.length > 0) {
+            try {
+                window.MathJax.typesetClear(mathElements)
+            } catch (error) {
+                console.warn('chatMath: failed to clear container MathJax state during dispose.', error)
+            }
+        }
+    }
+}