Streaming LLM Responses with Swift Concurrency

A chat UI that waits for the full LLM response feels broken, even when it’s fast. Streaming tokens as they’re generated is table stakes now. The good news: token streams and Swift Concurrency fit together almost suspiciously well. The bad news: the naive version has three bugs you won’t notice until users do.

The stream itself

An LLM generating tokens is an AsyncSequence that produces strings and can fail. AsyncThrowingStream is the bridge between the inference engine’s callback world and async/await:

func generate(prompt: String) -> AsyncThrowingStream<String, Error> {
    AsyncThrowingStream { continuation in
        let task = engine.startGeneration(prompt: prompt) { event in
            switch event {
            case .token(let text):  continuation.yield(text)
            case .finished:         continuation.finish()
            case .failed(let err):  continuation.finish(throwing: err)
            }
        }
        continuation.onTermination = { _ in
            task.cancel()   // <- the line everyone forgets
        }
    }
}

The onTermination handler is the first bug in most implementations. Without it, the stream and the generation have different lifetimes: the user leaves the screen, the for await loop ends, and the model keeps generating tokens into the void — burning battery and, on-device, holding megabytes of KV cache for a conversation nobody is reading.

Consuming it

On the view model side, the loop is almost embarrassingly simple:

@MainActor
final class ChatViewModel: ObservableObject {
    @Published var reply = ""
    private var generation: Task<Void, Never>?

    func send(_ prompt: String) {
        generation?.cancel()   // one generation at a time
        generation = Task {
            do {
                for try await token in llm.generate(prompt: prompt) {
                    reply += token
                }
            } catch is CancellationError {
                // user interrupted — keep the partial reply
            } catch {
                reply += "\n\n⚠️ Generation failed."
            }
        }
    }

    func stop() { generation?.cancel() }
}

Cancelling the Task ends the for await loop, which terminates the stream, which — thanks to onTermination — actually stops the engine. One cancellation path, top to bottom. When someone reports “the stop button doesn’t stop it,” it’s because one link in that chain is missing.

Note the CancellationError case: interruption is a feature in a chat app, not an error. The user tapped stop because the answer was already good enough, or already wrong. Keep the partial text.

The bug you’ll find in Instruments

Here’s the subtle one. A model can produce a great many tokens per second, and reply += token publishes a SwiftUI update for every one of them. Each update re-renders a growing Text, re-layouts the bubble, re-measures the scroll view. Early in a response nobody notices. Two thousand tokens into a long answer, the main thread is spending its entire budget on layout and the scroll starts hitching.

The fix isn’t clever, it’s a buffer — accumulate tokens and flush to the UI on a fixed cadence:

var buffer = ""
var lastFlush = ContinuousClock.now

for try await token in llm.generate(prompt: prompt) {
    buffer += token
    if ContinuousClock.now - lastFlush > .milliseconds(50) {
        reply += buffer
        buffer = ""
        lastFlush = .now
    }
}
reply += buffer   // whatever's left when the stream ends

Twenty UI updates a second is indistinguishable from per-token updates to a human eye, and it’s the difference between a smooth stream and a slideshow. This one change matters more than anything else in this post.

Keep the engine behind an actor

On-device inference is stateful — a loaded model, a KV cache, sometimes a single non-reentrant C API. Two overlapping generations will corrupt state or crash. An actor makes the exclusivity structural instead of disciplinary:

actor InferenceEngine {
    private var model: Model?

    func generate(prompt: String) -> AsyncThrowingStream<String, Error> {
        // callers queue here automatically
    }
}

Same argument as making illegal imports a compile error: rules enforced by the language don’t need code review to survive.

Test with a fake stream

The last win: everything above depends on AsyncThrowingStream<String, Error>, not on a real model. So tests inject a fake that yields scripted tokens with scripted delays — including a stream that fails halfway, and one that gets cancelled mid-word. Those two scenarios are exactly the ones you can’t reliably reproduce with a live model, and exactly the ones users hit daily.

The streaming layer is the part of an AI app you can get completely right before the model ever loads. Worth doing — it’s also the part users actually touch.