Enable speculative rewrite for streamed markdown snapshots

Examples/SwiftStreamingMarkdownSample/SwiftStreamingMarkdownSample/LLMChatInteractor.swift

@@ -53,7 +53,12 @@ final class LLMChatInteractor: ObservableObject {
       ) ?? markdown.endIndex
 
       let snapshot = String(markdown[..<endIndex])
-      let document = await parser.parse(text: snapshot, config: markdownConfig)
+      let isComplete = endIndex == markdown.endIndex
+      let document = await parser.parse(
+        text: snapshot,
+        config: markdownConfig,
+        speculativeRewrite: !isComplete
+      )
       await viewModel.updateAssistantMessage(id: messageID, document: document)
 
       if endIndex == markdown.endIndex { break }

README.md

@@ -236,9 +236,17 @@ struct ChatBubble: View {
 }
 ```
 
+Mid-stream snapshots are parsed with speculative rewriting enabled —
+half-typed emphasis and partial table headers are repaired so text doesn't
+jitter — and the final snapshot is re-parsed literally once the stream
+completes, so a deliberate trailing `*` is restored.
+
 If you'd rather drive `DocumentView` directly, parse each snapshot with
 `MarkdownParser.parse(text:config:)` and feed the resulting
-`RenderableDocument` into your view yourself.
+`RenderableDocument` into your view yourself. Pass
+`speculativeRewrite: true` for incomplete snapshots and keep the default
+`false` for the completed document — the same behavior
+`StreamedMarkdownView` applies internally.
 
 The bundled [sample app](Examples/SwiftStreamingMarkdownSample) demonstrates
 chunked streaming end-to-end with adjustable chunk size and interval, plus

Sources/MarkdownText/Parser/MarkdownParser.swift

@@ -30,11 +30,15 @@ extension MarkdownParser {
   /// - Parameters:
   ///   - text: The incoming text.
   ///   - config: Render configuration applied when building the renderable.
+  ///   - speculativeRewrite: Whether to speculatively rewrite incomplete
+  ///     trailing constructs (partial emphasis, partial table headers). Intended
+  ///     for incremental snapshots; complete documents should keep the default
+  ///     `false` so deliberate trailing markers render literally.
   /// - Returns: A `RenderableDocument` built from the parsed `Document`.
-  public func parse(text: String, config: MarkdownRenderConfig) async -> RenderableDocument {
+  public func parse(text: String, config: MarkdownRenderConfig, speculativeRewrite: Bool = false) async -> RenderableDocument {
     let document = await parse(
       text: text,
-      option: .init(speculativeRewrite: false, imageSupport: config.imageConfig.enabled)
+      option: .init(speculativeRewrite: speculativeRewrite, imageSupport: config.imageConfig.enabled)
     ).document
     return await RenderableDocument(document: document, config: config)
   }

Sources/MarkdownText/StreamedMarkdownView.swift

@@ -69,7 +69,8 @@ final class StreamedMarkdownController: ObservableObject {
 
   private let source: StreamedMarkdownSource
   private let parser = MarkdownParserImpl()
-  private var task: Task<Void, Never>?
+  /// The streaming task, exposed read-only so tests can await completion.
+  private(set) var task: Task<Void, Never>?
 
   init(
     source: StreamedMarkdownSource,
@@ -85,11 +86,26 @@ final class StreamedMarkdownController: ObservableObject {
     task?.cancel()
     task = Task { [weak self] in
       guard let self else { return }
+      var lastText: String?
       for await text in self.source.text {
         if Task.isCancelled { return }
-        let renderable = await self.parser.parse(text: text, config: self.config)
+        let renderable = await self.parser.parse(text: text, config: self.config, speculativeRewrite: true)
+        lastText = text
         if Task.isCancelled { return }
         await MainActor.run {
+          // Cancellation during the hop must not publish a stale frame.
+          guard !Task.isCancelled else { return }
+          self.markdownToRender = renderable
+        }
+      }
+      // The stream completed: re-render the final snapshot without speculative
+      // rewriting so deliberate trailing markers are not eaten.
+      if Task.isCancelled { return }
+      if let lastText {
+        let renderable = await self.parser.parse(text: lastText, config: self.config)
+        if Task.isCancelled { return }
+        await MainActor.run {
+          guard !Task.isCancelled else { return }
           self.markdownToRender = renderable
         }
       }

Tests/MarkdownTextTests/StreamedMarkdownControllerTests.swift

@@ -0,0 +1,95 @@
+//
+//  Copyright (c) Microsoft Corporation. All rights reserved.
+//  Licensed under the MIT License. See LICENSE in the project root for license information.
+//
+
+@testable import SwiftStreamingMarkdown
+import XCTest
+
+@MainActor
+final class StreamedMarkdownControllerTests: XCTestCase {
+
+  private final class StubSource: StreamedMarkdownSource {
+    let text: AsyncStream<String>
+
+    init(stream: AsyncStream<String>) {
+      self.text = stream
+    }
+  }
+
+  private func makeStream() -> (AsyncStream<String>, AsyncStream<String>.Continuation) {
+    return AsyncStream<String>.makeStream(of: String.self)
+  }
+
+  private func waitForRender(_ controller: StreamedMarkdownController) async {
+    for _ in 0..<200 where controller.markdownToRender.renderables.isEmpty {
+      try? await Task.sleep(ms: 5)
+    }
+  }
+
+  func test_partial_trailing_strong_is_speculatively_closed_midstream() async {
+    let (stream, continuation) = makeStream()
+    let controller = StreamedMarkdownController(source: StubSource(stream: stream), config: .default)
+    await controller.start()
+
+    continuation.yield("Yeah, this is **cool")
+    await waitForRender(controller)
+    XCTAssertEqual(controller.markdownToRender.plainText, "Yeah, this is cool")
+
+    continuation.yield("Yeah, this is **cooler**.")
+    continuation.finish()
+    await controller.task?.value
+    XCTAssertEqual(controller.markdownToRender.plainText, "Yeah, this is cooler.")
+    await controller.end()
+  }
+
+  func test_literal_trailing_marker_is_restored_when_stream_completes() async {
+    let (stream, continuation) = makeStream()
+    let controller = StreamedMarkdownController(source: StubSource(stream: stream), config: .default)
+    await controller.start()
+
+    // Mid-stream the trailing "*" is speculatively treated as emphasis.
+    continuation.yield("2 * 3")
+    await waitForRender(controller)
+    XCTAssertEqual(controller.markdownToRender.plainText, "2  3")
+
+    // Once the stream finishes the final snapshot re-renders literally.
+    continuation.finish()
+    await controller.task?.value
+    XCTAssertEqual(controller.markdownToRender.plainText, "2 * 3")
+    await controller.end()
+  }
+
+  func test_cancelled_controller_does_not_republish_on_stream_finish() async {
+    let (stream, continuation) = makeStream()
+    let controller = StreamedMarkdownController(source: StubSource(stream: stream), config: .default)
+    await controller.start()
+
+    continuation.yield("2 * 3")
+    await waitForRender(controller)
+    XCTAssertEqual(controller.markdownToRender.plainText, "2  3")
+
+    await controller.end()
+    continuation.finish()
+    try? await Task.sleep(ms: 50)
+    // The cancelled task must not publish the settled literal render.
+    XCTAssertEqual(controller.markdownToRender.plainText, "2  3")
+  }
+
+  func test_partial_table_header_is_hidden_until_complete() async {
+    let (stream, continuation) = makeStream()
+    let controller = StreamedMarkdownController(source: StubSource(stream: stream), config: .default)
+    await controller.start()
+
+    continuation.yield("intro\n\n| Month | Savings |")
+    await waitForRender(controller)
+    XCTAssertTrue(controller.markdownToRender.plainText.contains("intro"))
+    XCTAssertFalse(controller.markdownToRender.plainText.contains("Month"))
+
+    continuation.yield("intro\n\n| Month | Savings |\n| ----- | ------- |\n| Jan | 100 |")
+    continuation.finish()
+    await controller.task?.value
+    XCTAssertTrue(controller.markdownToRender.plainText.contains("Month"))
+    await controller.end()
+  }
+}