From dbcb2cc60c79bf2f8dac84db3505f1bbdfb69a51 Mon Sep 17 00:00:00 2001 From: Kaede Akino Date: Fri, 28 Aug 2026 15:32:35 +0800 Subject: [PATCH] Smooth agent streaming output --- .../Application/Agent/AgentTurnDriver.swift | 166 +++++++++------ .../Application/Agent/BalancedEmitter.swift | 102 +++++++++ .../Resources/LanguageModelChatUI.LICENSE.txt | 24 +++ TermindTests/BalancedEmitterTests.swift | 195 ++++++++++++++++++ TermindTests/ChatComposerStateTests.swift | 5 +- 5 files changed, 431 insertions(+), 61 deletions(-) create mode 100644 Termind/Application/Agent/BalancedEmitter.swift create mode 100644 Termind/Resources/LanguageModelChatUI.LICENSE.txt create mode 100644 TermindTests/BalancedEmitterTests.swift diff --git a/Termind/Application/Agent/AgentTurnDriver.swift b/Termind/Application/Agent/AgentTurnDriver.swift index 2ffa829f..621e4b75 100644 --- a/Termind/Application/Agent/AgentTurnDriver.swift +++ b/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 = [] 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) diff --git a/Termind/Application/Agent/BalancedEmitter.swift b/Termind/Application/Agent/BalancedEmitter.swift new file mode 100644 index 00000000..5b90e776 --- /dev/null +++ b/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? + 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 + } + } + } +} diff --git a/Termind/Resources/LanguageModelChatUI.LICENSE.txt b/Termind/Resources/LanguageModelChatUI.LICENSE.txt new file mode 100644 index 00000000..fbf85ddc --- /dev/null +++ b/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. diff --git a/TermindTests/BalancedEmitterTests.swift b/TermindTests/BalancedEmitterTests.swift new file mode 100644 index 00000000..eae5526b --- /dev/null +++ b/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? + + 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 { + 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) + } +} diff --git a/TermindTests/ChatComposerStateTests.swift b/TermindTests/ChatComposerStateTests.swift index 0d88bc56..a6e0f85c 100644 --- a/TermindTests/ChatComposerStateTests.swift +++ b/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)) } }