Smooth agent streaming output

Termind/Application/Agent/AgentTurnDriver.swift

@@ -46,74 +46,120 @@ nonisolated struct AgentTurnDriver: Sendable {
         // read as the model having searched four times for one thing.
         var recordedSearches: Set<String> = []
         channel.emit(.messageStarted(message))
+        let messageID = message.id
+
+        let reasoningEmitter = BalancedEmitter(
+            configuration: .init(duration: .seconds(1), frequency: 30)
+        ) { [channel] text in
+            channel.emit(.reasoningDelta(messageID: messageID, text: text))
+        }
+        let textEmitter = BalancedEmitter(
+            configuration: .init(duration: .milliseconds(500), frequency: 20)
+        ) { [channel] text in
+            channel.emit(.messageDelta(messageID: messageID, text: text))
+        }
+        var streamedCharacterCount = 0
 
         let request = AgentModelRequest(
             systemPrompt: context.systemPrompt,
             messages: context.messages,
             tools: context.tools
         )
-        for try await event in model.stream(request) {
-            try Task.checkCancellation()
-            switch event {
-            case .textDelta(let text):
-                message.text += text
-                channel.emit(.messageDelta(messageID: message.id, text: text))
-            case .textSnapshot(let text):
-                // The UI is reconciled by messageFinished after the provider
-                // stream closes, so replacing here cannot duplicate deltas.
-                message.text = text
-            case .reasoningDelta(let text):
-                if message.reasoning.isEmpty {
-                    message.reasoning.append(AgentReasoningBlock(
-                        text: text, createdAt: message.createdAt
-                    ))
-                } else {
-                    message.reasoning[message.reasoning.index(before: message.reasoning.endIndex)]
-                        .text += text
+        do {
+            for try await event in model.stream(request) {
+                try Task.checkCancellation()
+                switch event {
+                case .textDelta(let text):
+                    message.text += text
+                    try await reasoningEmitter.wait()
+                    if streamedCharacterCount >= 5_000 {
+                        await textEmitter.update(.init(duration: .seconds(1), frequency: 3))
+                    } else if streamedCharacterCount >= 2_000 {
+                        await textEmitter.update(.init(duration: .seconds(1), frequency: 9))
+                    } else if streamedCharacterCount >= 1_000 {
+                        await textEmitter.update(
+                            .init(duration: .milliseconds(500), frequency: 15)
+                        )
+                    }
+                    await textEmitter.add(text)
+                    streamedCharacterCount += text.count
+                case .textSnapshot(let text):
+                    // The UI is reconciled by messageFinished after the provider
+                    // stream closes, so replacing here cannot duplicate deltas.
+                    message.text = text
+                case .reasoningDelta(let text):
+                    if message.reasoning.isEmpty {
+                        message.reasoning.append(AgentReasoningBlock(
+                            text: text, createdAt: message.createdAt
+                        ))
+                    } else {
+                        message.reasoning[
+                            message.reasoning.index(before: message.reasoning.endIndex)
+                        ].text += text
+                    }
+                    try await textEmitter.wait()
+                    await reasoningEmitter.add(text)
+                case .reasoningSnapshot(let text):
+                    if message.reasoning.isEmpty {
+                        message.reasoning.append(AgentReasoningBlock(
+                            text: text, createdAt: message.createdAt
+                        ))
+                    } else {
+                        message.reasoning[
+                            message.reasoning.index(before: message.reasoning.endIndex)
+                        ].text = text
+                    }
+                case .toolCallDelta(let id, let name, let arguments):
+                    try await reasoningEmitter.wait()
+                    try await textEmitter.wait()
+                    if pending[id] == nil { order.append(id) }
+                    var entry = pending[id] ?? PendingToolCall()
+                    if let name { entry.name = name }
+                    entry.arguments += arguments
+                    pending[id] = entry
+                case .toolCallSnapshot(let id, let providerItemID, let name, let arguments):
+                    try await reasoningEmitter.wait()
+                    try await textEmitter.wait()
+                    if pending[id] == nil { order.append(id) }
+                    pending[id] = PendingToolCall(
+                        name: name, arguments: arguments,
+                        providerItemID: providerItemID
+                    )
+                case .providerItem(let item):
+                    if let id = item.objectValue?["id"]?.stringValue,
+                        let index = message.providerItems.firstIndex(where: {
+                            $0.objectValue?["id"]?.stringValue == id
+                        })
+                    {
+                        // `response.completed` can carry a richer final reasoning
+                        // item (notably encrypted_content) than output_item.done.
+                        message.providerItems[index] = item
+                    } else {
+                        message.providerItems.append(item)
+                    }
+                case .webSearch(let activity):
+                    try await reasoningEmitter.wait()
+                    try await textEmitter.wait()
+                    if recordedSearches.insert(activity.query).inserted {
+                        recordNativeWebSearch(activity, sourceMessageID: message.id)
+                    }
+                case .usage(let reported):
+                    usage = usage.merging(reported)
+                case .finished(let reason):
+                    stopReason = reason
                 }
-                channel.emit(.reasoningDelta(messageID: message.id, text: text))
-            case .reasoningSnapshot(let text):
-                if message.reasoning.isEmpty {
-                    message.reasoning.append(AgentReasoningBlock(
-                        text: text, createdAt: message.createdAt
-                    ))
-                } else {
-                    message.reasoning[message.reasoning.index(before: message.reasoning.endIndex)]
-                        .text = text
-                }
-            case .toolCallDelta(let id, let name, let arguments):
-                if pending[id] == nil { order.append(id) }
-                var entry = pending[id] ?? PendingToolCall()
-                if let name { entry.name = name }
-                entry.arguments += arguments
-                pending[id] = entry
-            case .toolCallSnapshot(let id, let providerItemID, let name, let arguments):
-                if pending[id] == nil { order.append(id) }
-                pending[id] = PendingToolCall(
-                    name: name, arguments: arguments,
-                    providerItemID: providerItemID
-                )
-            case .providerItem(let item):
-                if let id = item.objectValue?["id"]?.stringValue,
-                    let index = message.providerItems.firstIndex(where: {
-                        $0.objectValue?["id"]?.stringValue == id
-                    })
-                {
-                    // `response.completed` can carry a richer final reasoning
-                    // item (notably encrypted_content) than output_item.done.
-                    message.providerItems[index] = item
-                } else {
-                    message.providerItems.append(item)
-                }
-            case .webSearch(let activity):
-                if recordedSearches.insert(activity.query).inserted {
-                    recordNativeWebSearch(activity, sourceMessageID: message.id)
-                }
-            case .usage(let reported):
-                usage = usage.merging(reported)
-            case .finished(let reason):
-                stopReason = reason
             }
+            try await reasoningEmitter.wait()
+            try await textEmitter.wait()
+        } catch {
+            if Task.isCancelled {
+                await reasoningEmitter.cancel()
+                await textEmitter.cancel()
+            } else {
+                try? await reasoningEmitter.wait()
+                try? await textEmitter.wait()
+            }
+            throw error
         }
         message.toolCalls = try order.map { try assembled($0, pending: pending) }
         return Turn(message: message, reason: stopReason, usage: usage)

Termind/Application/Agent/BalancedEmitter.swift

@@ -0,0 +1,102 @@
+import Foundation
+
+/// Smooths a burst of streamed text into a bounded number of ordered updates.
+///
+/// Adapted from `BalancedEmitter` in LanguageModelChatUI by Lakr233. The
+/// original project is distributed under the MIT License; see
+/// `Resources/LanguageModelChatUI.LICENSE.txt`.
+actor BalancedEmitter {
+    nonisolated struct Configuration: Equatable, Sendable {
+        var duration: Duration
+        var frequency: Int
+
+        init(duration: Duration, frequency: Int) {
+            precondition(duration > .zero)
+            precondition(frequency > 0)
+            self.duration = duration
+            self.frequency = frequency
+        }
+    }
+
+    typealias Sleep = @Sendable (Duration) async throws -> Void
+    typealias Emit = @Sendable (String) async -> Void
+
+    private var buffer = ""
+    private var configuration: Configuration
+    private var batchSize = 1
+    private var drainTask: Task<Void, Never>?
+    private var generation = 0
+    private let sleep: Sleep
+    private let onEmit: Emit
+
+    init(
+        configuration: Configuration,
+        sleep: @escaping Sleep = { try await Task.sleep(for: $0) },
+        onEmit: @escaping Emit
+    ) {
+        self.configuration = configuration
+        self.sleep = sleep
+        self.onEmit = onEmit
+    }
+
+    func update(_ configuration: Configuration) {
+        self.configuration = configuration
+        recalculateBatchSize()
+    }
+
+    func add(_ chunk: String) {
+        guard !chunk.isEmpty else { return }
+        buffer += chunk
+        recalculateBatchSize()
+        dispatchLoopIfRequired()
+    }
+
+    func wait() async throws {
+        try await withTaskCancellationHandler {
+            await drainTask?.value
+            try Task.checkCancellation()
+        } onCancel: { [weak self] in
+            Task { await self?.cancel() }
+        }
+    }
+
+    func cancel() {
+        buffer = ""
+        generation += 1
+        drainTask?.cancel()
+        drainTask = nil
+    }
+
+    private func recalculateBatchSize() {
+        batchSize = max(1, Int(ceil(Double(buffer.count) / Double(configuration.frequency))))
+    }
+
+    private func dispatchLoopIfRequired() {
+        guard drainTask == nil else { return }
+        let generation = generation
+        drainTask = Task { [weak self] in
+            await self?.drain(generation: generation)
+        }
+    }
+
+    private func drain(generation: Int) async {
+        defer {
+            if self.generation == generation { drainTask = nil }
+        }
+
+        while !Task.isCancelled, self.generation == generation, !buffer.isEmpty {
+            let emitCount = min(batchSize, buffer.count)
+            let emitted = String(buffer.prefix(emitCount))
+            buffer.removeFirst(emitted.count)
+            await onEmit(emitted)
+
+            guard !buffer.isEmpty else { break }
+            let delay = configuration.duration / configuration.frequency
+            do {
+                try await sleep(delay)
+            } catch {
+                break
+            }
+        }
+    }
+}

Termind/Resources/LanguageModelChatUI.LICENSE.txt

@@ -0,0 +1,24 @@
+LanguageModelChatUI
+https://github.com/Lakr233/LanguageModelChatUI
+
+MIT License
+
+Copyright (c) 2026 @Lakr233
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

TermindTests/BalancedEmitterTests.swift

@@ -0,0 +1,195 @@
+import Foundation
+import Testing
+
+@testable import Termind
+
+@Suite
+struct BalancedEmitterTests {
+    @Test
+    func burstIsSplitIntoOrderedBatchesAndWaitDrainsIt() async throws {
+        let output = StringOutput()
+        let emitter = BalancedEmitter(
+            configuration: .init(duration: .seconds(1), frequency: 4),
+            sleep: { _ in await Task.yield() }
+        ) { text in
+            await output.append(text)
+        }
+
+        await emitter.add("abcdefgh")
+        try await emitter.wait()
+
+        #expect(await output.values == ["ab", "cd", "ef", "gh"])
+        #expect(await output.joined == "abcdefgh")
+    }
+
+    @Test
+    func updateChangesTheBatchingPolicy() async throws {
+        let output = StringOutput()
+        let emitter = BalancedEmitter(
+            configuration: .init(duration: .seconds(1), frequency: 8),
+            sleep: { _ in await Task.yield() }
+        ) { text in
+            await output.append(text)
+        }
+
+        await emitter.update(.init(duration: .seconds(1), frequency: 2))
+        await emitter.add("abcdefgh")
+        try await emitter.wait()
+
+        #expect(await output.values == ["abcd", "efgh"])
+    }
+
+    @Test
+    func emptyChunksAreIgnored() async throws {
+        let output = StringOutput()
+        let emitter = BalancedEmitter(
+            configuration: .init(duration: .seconds(1), frequency: 4),
+            sleep: { _ in await Task.yield() }
+        ) { text in
+            await output.append(text)
+        }
+
+        await emitter.add("")
+        try await emitter.wait()
+
+        #expect(await output.values.isEmpty)
+    }
+
+    @Test
+    func cancelDropsBufferedText() async {
+        let output = StringOutput()
+        let gate = EmitterSleepGate()
+        let emitter = BalancedEmitter(
+            configuration: .init(duration: .seconds(1), frequency: 4),
+            sleep: { _ in try await gate.sleep() }
+        ) { text in
+            await output.append(text)
+        }
+
+        await emitter.add("abcdefgh")
+        await gate.waitUntilSleeping()
+        await emitter.cancel()
+        await gate.release()
+        for _ in 0..<10 { await Task.yield() }
+
+        #expect(await output.values == ["ab"])
+    }
+
+    @Test
+    func cancellingAWaitCancelsTheDrain() async {
+        let output = StringOutput()
+        let emitter = BalancedEmitter(
+            configuration: .init(duration: .seconds(10), frequency: 2)
+        ) { text in
+            await output.append(text)
+        }
+
+        await emitter.add("abcd")
+        let waiter = Task { try await emitter.wait() }
+        await Task.yield()
+        waiter.cancel()
+
+        await #expect(throws: CancellationError.self) {
+            try await waiter.value
+        }
+        #expect(await output.values.count <= 1)
+    }
+
+    @Test
+    func turnDriverKeepsReasoningBeforeTextAndReturnsExactMessage() async throws {
+        let channel = AgentEventChannel(repository: InMemoryAgentRunRepository())
+        let events = AgentEventOutput()
+        let consumer = Task {
+            for await event in channel.events { await events.append(event) }
+        }
+        let driver = AgentTurnDriver(
+            model: BalancedScriptedModel(events: [
+                .reasoningDelta("think"), .textDelta("answer"), .finished(.completed),
+            ]),
+            tools: AgentToolRegistry([AnyAgentTool]()), channel: channel, runID: UUID()
+        )
+
+        let turn = try await driver.run(AgentModelContext(
+            systemPrompt: "system", messages: [], tools: []
+        ))
+        channel.emit(.messageFinished(turn.message))
+        channel.finish()
+        await consumer.value
+
+        let recorded = await events.values
+        let reasoning = recorded.compactMap { event -> String? in
+            guard case .reasoningDelta(_, let text) = event else { return nil }
+            return text
+        }.joined()
+        let answer = recorded.compactMap { event -> String? in
+            guard case .messageDelta(_, let text) = event else { return nil }
+            return text
+        }.joined()
+        let lastReasoning = recorded.lastIndex {
+            if case .reasoningDelta = $0 { true } else { false }
+        }
+        let firstAnswer = recorded.firstIndex {
+            if case .messageDelta = $0 { true } else { false }
+        }
+        let finished = recorded.firstIndex {
+            if case .messageFinished = $0 { true } else { false }
+        }
+
+        #expect(reasoning == "think")
+        #expect(answer == "answer")
+        #expect(lastReasoning != nil && firstAnswer != nil && lastReasoning! < firstAnswer!)
+        #expect(firstAnswer != nil && finished != nil && firstAnswer! < finished!)
+        #expect(turn.message.reasoning.map { $0.text }.joined() == "think")
+        #expect(turn.message.text == "answer")
+    }
+}
+
+private actor StringOutput {
+    private(set) var values: [String] = []
+    var joined: String { values.joined() }
+
+    func append(_ value: String) {
+        values.append(value)
+    }
+}
+
+private actor EmitterSleepGate {
+    private var isSleeping = false
+    private var continuation: CheckedContinuation<Void, Never>?
+
+    func sleep() async throws {
+        isSleeping = true
+        await withCheckedContinuation { continuation = $0 }
+        try Task.checkCancellation()
+    }
+
+    func waitUntilSleeping() async {
+        while !isSleeping { await Task.yield() }
+    }
+
+    func release() {
+        continuation?.resume()
+        continuation = nil
+    }
+}
+
+private struct BalancedScriptedModel: AgentModelStreaming {
+    let events: [AgentModelStreamEvent]
+
+    func stream(
+        _ request: AgentModelRequest
+    ) -> AsyncThrowingStream<AgentModelStreamEvent, any Error> {
+        AsyncThrowingStream { continuation in
+            for event in events { continuation.yield(event) }
+            continuation.finish()
+        }
+    }
+}
+
+private actor AgentEventOutput {
+    private(set) var values: [AgentEvent] = []
+
+    func append(_ value: AgentEvent) {
+        values.append(value)
+    }
+}

TermindTests/ChatComposerStateTests.swift

@@ -1678,7 +1678,10 @@ struct ChatComposerStateTests {
     }
 
     private func waitForTerminal(_ conversation: ConversationModel) async {
-        for _ in 0..<200 where !conversation.runState.isTerminal {
+        // Multi-turn scripted runs also exercise the bounded presentation
+        // drain, so their terminal event can legitimately arrive after the
+        // old one-second polling window.
+        for _ in 0..<1_000 where !conversation.runState.isTerminal {
             try? await Task.sleep(for: .milliseconds(5))
         }
     }