Add an implicit cancel() call in reset(), therefore avoiding potential programming errors.

This also help get rid of the drain() functions which no longer are necessary.

Author: stikves

swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift

@@ -42,6 +42,27 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable {
     private let generationTask = Mutex<Task<Void, Never>?>(nil)
     let config: ModelConfig
 
+    // MARK: - Lifecycle
+
+    var isBusy: Bool { engineInUse.load(ordering: .acquiring) }
+
+    func cancel() async throws {
+        let task = generationTask.withLock { v in
+            v?.cancel()
+            return v
+        }
+        await task?.value
+        await engine.computeStream.currentWorkCompleted()
+        generationTask.withLock { $0 = nil }
+    }
+
+    func reset() async throws {
+        try await cancel()
+        guard tryAcquireEngine() else { return }
+        defer { releaseEngine() }
+        engine.reset()
+    }
+
     init(
         config: ModelConfig,
         preparedModel: PreparedModel,
@@ -144,37 +165,6 @@ final class CoreAIPipelinedEngine: InferenceEngine, Sendable {
         return GenerationSequence(base: base, stopReasonStore: stopReasonStore)
     }
 
-    var isBusy: Bool { engineInUse.load(ordering: .acquiring) }
-
-    func cancel() async {
-        let task = generationTask.withLock { v in
-            v?.cancel()
-            return v
-        }
-        await task?.value
-        await engine.computeStream.currentWorkCompleted()
-        generationTask.withLock { $0 = nil }
-    }
-
-    /// Wait for any in-flight generate() Task to return the engine.
-    private func drain() {
-        var attempts = 0
-        while engineInUse.load(ordering: .acquiring) {
-            attempts += 1
-            if attempts > Self.drainTimeoutIterations {
-                fatalError("Engine not returned after drain() — tokenSequence Task stuck?")
-            }
-            Thread.sleep(forTimeInterval: 0.001)
-        }
-    }
-
-    func reset() {
-        drain()
-        guard tryAcquireEngine() else { return }
-        defer { releaseEngine() }
-        engine.reset()
-    }
-
     func cleanup() async throws {
         let cleanupSpan = InstrumentsProfiler.beginCleanup(engine: "CoreAI-Pipelined")
         if tryAcquireEngine() {

swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift

@@ -350,37 +350,23 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
 
     public var isBusy: Bool { generating.withLock { $0 } }
 
-    public func cancel() async {
+    public func cancel() async throws {
         cancelRequested.withLock { $0 = true }
         var attempts = 0
         while generating.withLock({ $0 }) {
             attempts += 1
             if attempts > Self.drainTimeoutIterations {
-                CLILogger.log(
-                    "WARNING: Sequential engine cancel() drain timeout after \(Self.drainTimeoutIterations) iterations",
-                    component: "CoreAISequential"
-                )
-                break
+                cancelRequested.withLock { $0 = false }
+                throw InferenceRuntimeError.invalidState(
+                    "Sequential engine cancel() timeout — generation Task stuck?")
             }
             await Task.yield()
         }
         cancelRequested.withLock { $0 = false }
     }
 
-    /// Wait for any in-flight generate() Task to finish.
-    private func drain() {
-        var attempts = 0
-        while generating.withLock({ $0 }) {
-            attempts += 1
-            if attempts > Self.drainTimeoutIterations {
-                fatalError("Sequential engine drain() timeout — generation Task stuck?")
-            }
-            Thread.sleep(forTimeInterval: 0.001)
-        }
-    }
-
-    public func reset() {
-        drain()
+    public func reset() async throws {
+        try await cancel()
         let resetSpan = InstrumentsProfiler.beginReset(engine: "CoreAIClean")
         processedTokenCount = 0
         zeroFill(&keyCache)

swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIStaticShapeEngine.swift

@@ -538,37 +538,23 @@ public final class StaticShapeEngine: InferenceEngine, @unchecked Sendable {
 
     public var isBusy: Bool { generating.withLock { $0 } }
 
-    public func cancel() async {
+    public func cancel() async throws {
         cancelRequested.withLock { $0 = true }
         var attempts = 0
         while generating.withLock({ $0 }) {
             attempts += 1
             if attempts > Self.drainTimeoutIterations {
-                CLILogger.log(
-                    "WARNING: StaticShape engine cancel() drain timeout after \(Self.drainTimeoutIterations) iterations",
-                    component: "StaticShape"
-                )
-                break
+                cancelRequested.withLock { $0 = false }
+                throw InferenceRuntimeError.invalidState(
+                    "StaticShape engine cancel() timeout — generation stuck?")
             }
             await Task.yield()
         }
         cancelRequested.withLock { $0 = false }
     }
 
-    /// Wait for any in-flight generate() iteration to finish.
-    private func drain() {
-        var attempts = 0
-        while generating.withLock({ $0 }) {
-            attempts += 1
-            if attempts > Self.drainTimeoutIterations {
-                fatalError("StaticShape engine drain() timeout — generation stuck?")
-            }
-            Thread.sleep(forTimeInterval: 0.001)
-        }
-    }
-
-    public func reset() {
-        drain()
+    public func reset() async throws {
+        try await cancel()
         let resetSpan = InstrumentsProfiler.beginReset(engine: "StaticShape")
         processedTokenCount = 0
         resetSpan.end()
@@ -578,7 +564,7 @@ public final class StaticShapeEngine: InferenceEngine, @unchecked Sendable {
         for fnName in extendFunctionNames {
             self.functions[fnName] = try Self.requireFunction(model: model, functionName: fnName)
         }
-        reset()
+        try await reset()
     }
 }
 

swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift

@@ -110,9 +110,8 @@ public protocol InferenceEngine: Sendable {
     var isBusy: Bool { get }
 
     /// Cancel in-flight generation and wait for the engine to become idle.
-    ///
-    /// Best-effort: logs a warning if the drain times out but does not trap.
-    func cancel() async
+    /// Throws if the engine doesn't become idle within the drain timeout.
+    func cancel() async throws
 
     /// Reset internal state including KV cache.
     func reset() async throws
@@ -166,15 +165,14 @@ extension InferenceConfiguration {
 
 extension InferenceEngine {
     /// Maximum iterations for drain loops before timing out.
-    /// - `cancel()` logs a warning and breaks (best-effort).
-    /// - `reset()` traps (hard lifecycle boundary).
+    /// Traps if the engine doesn't become idle within this window.
     static var drainTimeoutIterations: Int { 5000 }
 
     /// Default: engine is not busy.
     public var isBusy: Bool { false }
 
     /// Default no-op cancel.
-    public func cancel() async {}
+    public func cancel() async throws {}
 }
 
 extension InferenceEngine {

swift/Sources/CoreAILanguageModels/LanguageModel/CoreAILanguageModel.swift

@@ -290,7 +290,6 @@ public struct CoreAILanguageModel: LanguageModel {
             let maxTokens = request.generationOptions.maximumResponseTokens ?? 512
 
             // Reset engine state for new generation
-            await engine.cancel()
             try await engine.reset()
 
             // FoundationModels now threads entry identity itself based on event

swift/Tests/LanguageModelsTests/TestUtilities.swift

@@ -160,7 +160,7 @@ class MockEngine: InferenceEngine, @unchecked Sendable {
 
     var isBusy: Bool { false }
 
-    func cancel() async {}
+    func cancel() async throws {}
 
     func reset() async throws {
         resetCalled = true

swift/Tests/LanguageModelsTests/UnifiedGenerationAPITests.swift

@@ -302,7 +302,7 @@ struct CancelAPITests {
     @Test("MockEngine cancel() is a no-op")
     func mockEngineCancelNoOp() async {
         let engine = MockEngine(tokens: [1, 2, 3])
-        await engine.cancel()
+        try await engine.cancel()
         #expect(engine.isBusy == false)
     }
 
@@ -312,9 +312,9 @@ struct CancelAPITests {
     }
 
     @Test("cancel() before generate does not crash")
-    func cancelBeforeGenerate() async {
+    func cancelBeforeGenerate() async throws {
         let engine = MockEngine(tokens: [1, 2, 3])
-        await engine.cancel()
+        try await engine.cancel()
         // Should be able to generate normally after cancel
         var tokens: [Int32] = []
         for try await output in try engine.generate(