From e3bf3b20adb0c466d6401547505ef38c09ae8e18 Mon Sep 17 00:00:00 2001
From: John Rocky <john-rocky@users.noreply.github.com>
Date: Thu, 30 Apr 2026 08:48:06 +0900
Subject: [PATCH] feat(runtime): mid-stream cancellation for AsyncStream
 generation
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

AsyncStream does NOT auto-cancel the producer Task when the consumer
breaks out of `for await`. A user who taps stop or navigates away
during a long generation kept burning ANE energy until the loop hit
maxDecode. This wires the missing piece:

- bind the producer Task to a `genTask` handle
- check Task.checkCancellation() at chunk boundaries in both prefill
  paths (hybrid + decode-loop fallback) and at the top of the decode
  loop, so cancellation lands within one chunk / one token
- catch CancellationError silently — it's the expected user path
- continuation.onTermination → genTask.cancel() so a stream deinit or
  the consumer dropping subscription stops generation immediately

Net effect: no API change, no token-output difference, but cancelling
a generation now releases the ANE within one decode step (~30 ms)
instead of seconds-to-minutes. Extracted from feat/litert-perf-adoptions
(commit 8598388 §T4); other items in that branch will land separately.
---
 Sources/CoreMLLLM/CoreMLLLM.swift | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/Sources/CoreMLLLM/CoreMLLLM.swift b/Sources/CoreMLLLM/CoreMLLLM.swift
index cd4f295..90cfb5d 100644
--- a/Sources/CoreMLLLM/CoreMLLLM.swift
+++ b/Sources/CoreMLLLM/CoreMLLLM.swift
@@ -1003,7 +1003,13 @@ public final class CoreMLLLM: @unchecked Sendable {
         }
 
         return AsyncStream { continuation in
-            Task(priority: decodePriority) {
+            // T4 cancellation: AsyncStream does NOT auto-cancel the producer
+            // task when the consumer breaks out of `for await`. Bind a
+            // handle so onTermination can call cancel(). Combined with
+            // Task.checkCancellation() at chunk and decode-step boundaries,
+            // a user-initiated cancel (or stream deinit) stops generation
+            // promptly instead of running to maxDecode.
+            let genTask = Task(priority: decodePriority) {
                 do {
                     let IMAGE_TOKEN_ID = 258880
                     let AUDIO_TOKEN_ID = 258881
@@ -1077,6 +1083,7 @@ public final class CoreMLLLM: @unchecked Sendable {
                             engine.currentPosition = prefillLen
 
                             for step in prefillLen..<tokens.count {
+                                try Task.checkCancellation()
                                 let tid = tokens[step]
                                 try autoreleasepool {
                                     if let emb = multimodalEmbedding(for: tid) {
@@ -1090,6 +1097,7 @@ public final class CoreMLLLM: @unchecked Sendable {
                             }
                         } else {
                             for (step, tid) in tokens.enumerated() {
+                                try Task.checkCancellation()
                                 try autoreleasepool {
                                     if let emb = multimodalEmbedding(for: tid) {
                                         nextID = try engine.predictStep(tokenID: 0, position: step,
@@ -1161,6 +1169,10 @@ public final class CoreMLLLM: @unchecked Sendable {
                         decodeLoop: while tokenCount < maxDecode {
                             if eosIDs.contains(Int(nid)) { break }
                             if engine.currentPosition >= ctxLimit { break }
+                            // T4: stop the loop the moment the consumer
+                            // unsubscribes — wired via continuation.onTermination
+                            // → genTask.cancel() below.
+                            try Task.checkCancellation()
 
                             let useEagle = (eagleSpec != nil) && didFirstDecode
                                 && engine.canSpeculate
@@ -1343,11 +1355,16 @@ public final class CoreMLLLM: @unchecked Sendable {
                             pos += 1
                         }
                     }
+                } catch is CancellationError {
+                    // T4: user-cancelled. No log spam — this is expected.
                 } catch {
                     print("[CoreMLLLM] Error: \(error)")
                 }
                 continuation.finish()
             }
+            continuation.onTermination = { @Sendable _ in
+                genTask.cancel()
+            }
         }
     }