Keep a streaming reply as plain text until it is final (1.2.2)

CHANGELOG.md

@@ -1,5 +1,14 @@
 # Changelog
 
+## 1.2.2
+
+- Fixes the chat window freezing during a long reply (any model) and then showing the rest at once, and the transcript going blank when a reply finished. Sampled on an M2 Pro, the main thread was 100% busy in SwiftUI/AppKit layout while the engine (which streams at 6–7 ms per token) had already finished the reply into the stream buffer. Five causes, all in the chat view:
+  - Every paragraph of the streaming reply was a selectable, AppKit-backed text view that SwiftUI re-created and hit-tested (hover tracking) on each update. A streaming reply is plain text and becomes selectable once it is final; rows have no hover tracking while a reply streams.
+  - The live-stats dot, the "Thinking…" shimmer and the typing dots were repeat-forever animations, rendering 60 frames a second, each frame with a full AppKit layout and hit-test pass. They are stepped now.
+  - Every token mutated the conversation store, so the sidebar, the inspector (with its AppKit sliders) and the whole chat re-rendered per update. The reply streams into live state that only the in-progress row reads; the conversation is written once at the end.
+  - Finished paragraphs keep their view; only the paragraph being written re-renders.
+  - Updates are paced by their measured cost on the main thread (at most 20 a second, never more often than twice what the previous one took), so the UI keeps most of its time whatever the transcript size.
+  - The message list is a plain stack: a lazy one could be left scrolled past its stale content when the reply's layout changed at the end, which showed an empty transcript.
 ## 1.2.1
 
 - Fixes the GPU out-of-memory error when running a large model such as Gemma 4 12B on a 16 GB Mac. MLX writes `safetensors` without the 16-byte alignment the kernels need, so every tensor was copied into its own Metal buffer while the mapped file stayed referenced: the GPU held two copies of the weights (18 GB for a 9 GB model). Shards are now copied once, in file order, into a page-aligned buffer and the mapping is released. Loading is faster too: Gemma 4 12B 156 s → 25 s from disk, 11 s from the page cache; Mistral 7B 3.3 → 2.9 s.

Sources/MLXEngine/Version.swift

@@ -1,7 +1,7 @@
 /// The MLX Runner release version. The release workflow refuses a tag that doesn't match it,
 /// and scripts/bundle.sh writes it into the app's Info.plist.
 public enum MLXRunnerVersion {
-    public static let current = "1.2.1"
+    public static let current = "1.2.2"
     public static let repository = "BackendStack21/mlx-runner"
     public static let homepage = "https://mlx-runner.21no.de"
 }

Sources/MLXRunner/AppState.swift

@@ -112,6 +112,14 @@ final class AppState {
     var draft: String = ""
     var liveTokens = 0
     var liveTokensPerSecond: Double = 0
+    /// The reply being generated. Streamed here, not into `conversations`: every view that
+    /// reads `conversations` (sidebar, inspector with its AppKit sliders, the whole chat) would
+    /// otherwise re-render on every token, which froze the window for long replies. Only the
+    /// in-progress message row reads these; the final text is written to the conversation once.
+    var liveReplyID: UUID?
+    var liveContent = ""
+    var liveReasoning: String?
+    var liveReasoningSeconds: Double?
     /// Non-fatal persistence problem to show the user (e.g. an unreadable chats.json).
     var persistenceError: String?
 
@@ -347,12 +355,19 @@ final class AppState {
             var raw = ""
             var pending = ""
             var lastFlush = Date()
+            // What the previous update cost the main thread (state change through layout and
+            // display): the next one waits twice that, so the UI keeps most of its time.
+            var lastCost: TimeInterval = 0
             var thinkStart: Date?
             var thinkEnd: Date?
             var genStart: Date?
             var tokenCount = 0
             liveTokens = 0
             liveTokensPerSecond = 0
+            liveReplyID = replyID
+            liveContent = ""
+            liveReasoning = nil
+            liveReasoningSeconds = nil
             @MainActor func flush(_ final: Bool = false) {
                 raw += pending
                 pending = ""
@@ -364,10 +379,17 @@ final class AppState {
                     liveTokens = tokenCount
                     liveTokensPerSecond = Double(tokenCount) / max(Date().timeIntervalSince(g), 0.001)
                 }
-                update(id, replyID) { m in
-                    m.content = final ? parsed.content.trimmingCharacters(in: .whitespacesAndNewlines) : parsed.content
-                    m.reasoning = parsed.reasoning
-                    if let s = thinkStart { m.reasoningSeconds = (thinkEnd ?? Date()).timeIntervalSince(s) }
+                if final {
+                    update(id, replyID) { m in
+                        m.content = parsed.content.trimmingCharacters(in: .whitespacesAndNewlines)
+                        m.reasoning = parsed.reasoning
+                        if let s = thinkStart { m.reasoningSeconds = (thinkEnd ?? Date()).timeIntervalSince(s) }
+                    }
+                    liveReplyID = nil
+                } else {
+                    liveContent = parsed.content
+                    liveReasoning = parsed.reasoning
+                    if let s = thinkStart { liveReasoningSeconds = (thinkEnd ?? Date()).timeIntervalSince(s) }
                 }
             }
             do {
@@ -380,7 +402,16 @@ final class AppState {
                     case .text(let t):
                         pending += t
                         tokenCount += 1
-                        if Date().timeIntervalSince(lastFlush) > 0.033 { flush() }
+                        // Each update re-lays out and redraws the chat (tens of ms for a long reply),
+                        // so updates are paced by their measured cost: at most 20 a second, and
+                        // never more often than twice what the last one took.
+                        let interval = min(0.5, max(0.05, lastCost * 2))
+                        if Date().timeIntervalSince(lastFlush) > interval {
+                            let started = Date()
+                            flush()
+                            // Runs on the next run-loop turn, after this update's layout and display.
+                            DispatchQueue.main.async { lastCost = Date().timeIntervalSince(started) }
+                        }
                     case .finished(let s):
                         flush(true)
                         update(id, replyID) { $0.stats = s }

Sources/MLXRunner/ChatView.swift

@@ -216,7 +216,9 @@ struct MessageList: View {
     var body: some View {
         ScrollViewReader { proxy in
             ScrollView {
-                LazyVStack(spacing: 22) {
+                // A plain stack: a lazy one could be left scrolled past its (stale) content when a
+                // reply's layout changed at the end of generation, showing an empty transcript.
+                VStack(spacing: 22) {
                     Color.clear.frame(height: 8)
                     ForEach(conversation.messages) { m in
                         MessageRow(message: m, conversationID: conversation.id,
@@ -231,9 +233,13 @@ struct MessageList: View {
             }
             .scrollContentBackground(.hidden)
             .defaultScrollAnchor(.bottom)
-            .onChange(of: conversation.messages.last?.content) { _, _ in
+            // The streaming reply lives in the live state (not in the message) until it is final.
+            .onChange(of: state.liveContent.count) { _, _ in
                 if state.generatingConversation == conversation.id { proxy.scrollTo("bottom", anchor: .bottom) }
             }
+            .onChange(of: conversation.messages.last?.content) { _, _ in
+                proxy.scrollTo("bottom", anchor: .bottom)
+            }
             .onChange(of: conversation.messages.count) { _, _ in
                 withAnimation(.snappy) { proxy.scrollTo("bottom", anchor: .bottom) }
             }
@@ -291,6 +297,11 @@ struct MessageRow: View {
             }
             .onHover { hovering = $0 }
         } else {
+            // While this reply streams, its text lives in the state's live properties.
+            let live = generating && state.liveReplyID == message.id
+            let content = live ? state.liveContent : message.content
+            let reasoning = live ? state.liveReasoning : message.reasoning
+            let reasoningSeconds = live ? state.liveReasoningSeconds : message.reasoningSeconds
             HStack(alignment: .top, spacing: 12) {
                 Image(systemName: "sparkles")
                     .font(.system(size: 13, weight: .bold))
@@ -299,9 +310,9 @@ struct MessageRow: View {
                     .background(Theme.gradient, in: Circle())
                     .shadow(color: Theme.accent.opacity(0.4), radius: 6, y: 2)
                 VStack(alignment: .leading, spacing: 10) {
-                    if let r = message.reasoning, !r.isEmpty {
-                        ReasoningView(text: r, seconds: message.reasoningSeconds,
-                                      inProgress: generating && message.content.isEmpty)
+                    if let r = reasoning, !r.isEmpty {
+                        ReasoningView(text: r, seconds: reasoningSeconds,
+                                      inProgress: generating && content.isEmpty)
                     }
                     if let err = message.error {
                         Label(err, systemImage: "exclamationmark.triangle.fill")
@@ -310,23 +321,32 @@ struct MessageRow: View {
                             .padding(10)
                             .background(.orange.opacity(0.1), in: RoundedRectangle(cornerRadius: 10))
                     }
-                    if message.content.isEmpty && message.reasoning == nil && generating {
+                    if content.isEmpty && reasoning == nil && generating {
                         TypingIndicator()
-                    } else if !message.content.isEmpty {
-                        MarkdownView(text: message.content)
+                    } else if !content.isEmpty {
+                        if generating {
+                            // Finished blocks keep their view; only the block being written re-renders.
+                            let parts = MarkdownParser.splitStreaming(content)
+                            if !parts.head.isEmpty { MarkdownView(text: parts.head, selectable: false) }
+                            if !parts.tail.isEmpty { MarkdownView(text: parts.tail, selectable: false) }
+                        } else {
+                            MarkdownView(text: content)
+                        }
                     }
                     if !generating {
                         footer
                     } else if state.liveTokens > 0 {
                         LiveStats(tps: state.liveTokensPerSecond, tokens: state.liveTokens)
-                    } else if let s = message.stats, s.promptTokens > 0, message.content.isEmpty, message.reasoning == nil {
+                    } else if let s = message.stats, s.promptTokens > 0, content.isEmpty, reasoning == nil {
                         Text("Processed \(s.promptTokens) prompt tokens in \(String(format: "%.2f", s.timeToFirstToken))s")
                             .font(.system(size: 11)).foregroundStyle(.tertiary)
                     }
                 }
                 Spacer(minLength: 40)
             }
-            .onHover { hovering = $0 }
+            // No hover tracking while the reply streams: a tracking area makes AppKit
+            // re-hit-test the whole window on every update.
+            .onHover { if !generating { hovering = $0 } }
         }
     }
 
@@ -412,10 +432,13 @@ struct ReasoningView: View {
             }
             .buttonStyle(.plain)
             if expanded || inProgress {
-                Text(inProgress ? String(text.suffix(1200)) : text)
+                let body = Text(inProgress ? String(text.suffix(1200)) : text)
                     .font(.system(size: 12.5))
                     .foregroundStyle(.secondary)
-                    .textSelection(.enabled)
+                // Selectable once final (see MarkdownView.selectable).
+                Group {
+                    if inProgress { body } else { body.textSelection(.enabled) }
+                }
                     .padding(.leading, 12)
                     .overlay(alignment: .leading) {
                         Rectangle().fill(Theme.gradient).frame(width: 2).opacity(0.6)
@@ -435,22 +458,18 @@ struct TypingIndicator: View {
             }
         }
         .padding(.vertical, 6)
-        .onAppear { withAnimation(.linear(duration: 1.2).repeatForever(autoreverses: false)) { phase = .pi } }
+        .task { while !Task.isCancelled { try? await Task.sleep(nanoseconds: 120_000_000); phase += 0.45 } }   // stepped, see LiveStats
     }
 }
 
+/// A slow blink. (A moving-gradient shimmer was a repeat-forever animation: it rendered every
+/// frame and, on this window, dragged an AppKit layout and hit-test pass with each one.)
 struct Shimmer: ViewModifier {
-    @State private var x: CGFloat = -1
+    @State private var on = true
     func body(content: Content) -> some View {
         content
-            .overlay(
-                LinearGradient(colors: [.clear, .white.opacity(0.7), .clear], startPoint: .leading, endPoint: .trailing)
-                    .scaleEffect(x: 0.5)
-                    .offset(x: x * 80)
-                    .blendMode(.plusLighter)
-                    .mask(content)
-            )
-            .onAppear { withAnimation(.linear(duration: 1.4).repeatForever(autoreverses: false)) { x = 1 } }
+            .opacity(on ? 1 : 0.55)
+            .task { while !Task.isCancelled { try? await Task.sleep(nanoseconds: 600_000_000); on.toggle() } }
     }
 }
 
@@ -627,17 +646,17 @@ struct LiveStats: View {
     let tokens: Int
     @State private var pulse = false
     var body: some View {
+        // Stepped, not animated: a repeat-forever animation renders every frame, and on this
+        // window each frame costs an AppKit layout and hit-test pass (the chat froze for it).
         HStack(spacing: 8) {
             Circle().fill(Theme.gradient).frame(width: 7, height: 7)
-                .scaleEffect(pulse ? 1.25 : 0.8)
-                .opacity(pulse ? 1 : 0.6)
-                .animation(.easeInOut(duration: 0.7).repeatForever(), value: pulse)
+                .opacity(pulse ? 1 : 0.5)
             Text(String(format: "%.0f tok/s", tps))
                 .font(.system(size: 11, weight: .semibold).monospacedDigit())
                 .foregroundStyle(Theme.gradient)
                 .contentTransition(.numericText())
             Text("\(tokens) tokens").font(.system(size: 11).monospacedDigit()).foregroundStyle(.tertiary)
         }
-        .onAppear { pulse = true }
+        .task { while !Task.isCancelled { try? await Task.sleep(nanoseconds: 700_000_000); pulse.toggle() } }
     }
 }

Sources/MLXRunner/MarkdownView.swift

@@ -7,16 +7,25 @@ import SwiftUI
 struct MarkdownView: View {
     let text: String
     var baseSize: CGFloat = 14
+    /// Text selection makes every paragraph an AppKit-backed view that SwiftUI re-creates and
+    /// hit-tests on each update; a reply that is still streaming (re-rendered 20 times a
+    /// second) stays plain SwiftUI text and becomes selectable once it is final. With
+    /// selection on while streaming, a long reply froze the main thread for seconds.
+    var selectable = true
 
     var body: some View {
         let blocks = MarkdownParser.parse(text)
-        VStack(alignment: .leading, spacing: 10) {
+        let stack = VStack(alignment: .leading, spacing: 10) {
             ForEach(Array(blocks.enumerated()), id: \.offset) { _, b in
                 view(for: b)
             }
         }
-        .textSelection(.enabled)
         .frame(maxWidth: .infinity, alignment: .leading)
+        if selectable {
+            stack.textSelection(.enabled)
+        } else {
+            stack
+        }
     }
 
     @ViewBuilder
@@ -200,6 +209,30 @@ enum MDBlock {
 }
 
 enum MarkdownParser {
+    /// Splits streaming text into the blocks that are complete (up to the last blank line
+    /// outside a code fence) and the block still being written. The head changes only when a
+    /// new block starts, so a view keyed on it is not re-evaluated on every token; the tail is
+    /// the only part re-rendered per update.
+    static func splitStreaming(_ text: String) -> (head: String, tail: String) {
+        var fences = 0
+        var cut: String.Index?
+        var i = text.startIndex
+        var lineStart = true
+        while i < text.endIndex {
+            if lineStart, text[i...].hasPrefix("```") { fences += 1 }
+            if text[i] == "\n" {
+                let next = text.index(after: i)
+                if next < text.endIndex, text[next] == "\n", fences % 2 == 0 { cut = text.index(after: next) }
+                lineStart = true
+            } else {
+                lineStart = false
+            }
+            i = text.index(after: i)
+        }
+        guard let c = cut else { return ("", text) }
+        return (String(text[..<c]), String(text[c...]))
+    }
+
     static func parse(_ text: String) -> [MDBlock] {
         var blocks: [MDBlock] = []
         let lines = text.components(separatedBy: "\n")