Stop pipelined generation when consumer drops the stream

swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift

@@ -45,6 +45,12 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable {
     typealias ConfigType = ModelConfig
 
     nonisolated(unsafe) private var engine: EngineImpl
+
+    /// Token history for implicit prefix caching. Marked nonisolated(unsafe) because
+    /// mutations are serialized by the generation lifecycle: generate() awaits any prior
+    /// Task before starting, and the forwarding `async let` only appends tokens while
+    /// runCompletion holds the engine lock. No concurrent writes are possible when the
+    /// cancel-and-await contract is upheld.
     nonisolated(unsafe) private var history = TokenHistory()
     nonisolated(unsafe) private(set) var lastPrefixHitCount: Int = 0
     private let engineInUse = Atomic<Bool>(false)
@@ -147,6 +153,10 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable {
                 let (tokenStream, tokenContinuation) =
                     AsyncThrowingStream<InferenceEngine.TokenId, any Error>.makeStream()
 
+                outputContinuation.onTermination = { @Sendable _ in
+                    tokenContinuation.finish()
+                }
+
                 // Implicit prefix caching: resolve input against history
                 var (commonPrefix, resolvedNewTokens) = self.history.resolve(input: input)
                 self.lastPrefixHitCount = commonPrefix
@@ -181,9 +191,12 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable {
                 async let forwarding: Void = {
                     do {
                         for try await token in tokenStream {
-                            // Track generated tokens in history
                             self.history.append(token)
-                            outputContinuation.yield(InferenceOutput(tokenId: token))
+                            let result = outputContinuation.yield(InferenceOutput(tokenId: token))
+                            if case .terminated = result {
+                                tokenContinuation.finish()
+                                break
+                            }
                         }
                     } catch {
                         outputContinuation.finish(throwing: error)

swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift

@@ -349,7 +349,12 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
         samplingConfiguration: SamplingConfiguration,
         inferenceOptions: InferenceOptions
     ) async throws -> GenerationSequence {
-        // Implicit prefix caching: resolve before creating Iterator.
+        // Cancel any prior generation so its Iterator stops on next poll.
+        _activeToken.withLock {
+            $0?.cancel()
+            $0 = nil
+        }
+
         // Implicit prefix caching: resolve input against history.
         if history.count > 0 {
             let (commonPrefix, _) = history.resolve(input: input)

swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift

@@ -775,6 +775,10 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec
         samplingConfiguration: SamplingConfiguration,
         inferenceOptions: InferenceOptions
     ) async throws -> GenerationSequence {
+        _activeToken.withLock {
+            $0?.cancel()
+            $0 = nil
+        }
         let token = GenerationToken()
         _activeToken.withLock { $0 = token }
         return GenerationSequence(
@@ -800,6 +804,10 @@ public final class CoreAISequentialVLMEngine: MultimodalInferenceEngine, @unchec
         samplingConfiguration: SamplingConfiguration,
         inferenceOptions: InferenceOptions
     ) async throws -> GenerationSequence {
+        _activeToken.withLock {
+            $0?.cancel()
+            $0 = nil
+        }
         let token = GenerationToken()
         _activeToken.withLock { $0 = token }
         return GenerationSequence(

swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift

@@ -332,6 +332,12 @@ public final class StaticShapeEngine: InferenceEngine, @unchecked Sendable {
         samplingConfiguration: SamplingConfiguration,
         inferenceOptions: InferenceOptions
     ) async throws -> GenerationSequence {
+        // Cancel any prior generation so its Iterator stops on next poll.
+        _activeToken.withLock {
+            $0?.cancel()
+            $0 = nil
+        }
+
         // Implicit prefix caching: resolve input against history.
         if history.count > 0 {
             let (commonPrefix, _) = history.resolve(input: input)