diff --git a/Examples/SwiftStreamingMarkdownSample/SwiftStreamingMarkdownSample/LLMChatInteractor.swift b/Examples/SwiftStreamingMarkdownSample/SwiftStreamingMarkdownSample/LLMChatInteractor.swift index d5f8c34c..2a482a5d 100644 --- a/Examples/SwiftStreamingMarkdownSample/SwiftStreamingMarkdownSample/LLMChatInteractor.swift +++ b/Examples/SwiftStreamingMarkdownSample/SwiftStreamingMarkdownSample/LLMChatInteractor.swift @@ -53,7 +53,12 @@ final class LLMChatInteractor: ObservableObject { ) ?? markdown.endIndex let snapshot = String(markdown[.. 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) } diff --git a/Sources/MarkdownText/StreamedMarkdownView.swift b/Sources/MarkdownText/StreamedMarkdownView.swift index f55e82c1..fd2ade5e 100644 --- a/Sources/MarkdownText/StreamedMarkdownView.swift +++ b/Sources/MarkdownText/StreamedMarkdownView.swift @@ -69,7 +69,8 @@ final class StreamedMarkdownController: ObservableObject { private let source: StreamedMarkdownSource private let parser = MarkdownParserImpl() - private var task: Task? + /// The streaming task, exposed read-only so tests can await completion. + private(set) var task: Task? 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 } } diff --git a/Tests/MarkdownTextTests/StreamedMarkdownControllerTests.swift b/Tests/MarkdownTextTests/StreamedMarkdownControllerTests.swift new file mode 100644 index 00000000..b32ed33f --- /dev/null +++ b/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 + + init(stream: AsyncStream) { + self.text = stream + } + } + + private func makeStream() -> (AsyncStream, AsyncStream.Continuation) { + return AsyncStream.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() + } +}