From eaa11743c9b490b8bca43326479de3bdb9c2bb21 Mon Sep 17 00:00:00 2001 From: james-333i Date: Thu, 27 Aug 2026 16:02:53 -0700 Subject: [PATCH] Cancel the generation when a response stream's consumer stops early MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit streamResponse wraps the model's stream in a relay that appends the response to the transcript when generation completes. The relay had no termination handler, so a consumer that stopped iterating early — a break, a cancelled task, or a dropped stream — left the relay consuming and the backend generating to its token limit in the background. That ghost generation kept the model and its growing KV cache alive for tens of seconds after the response visibly ended. Cancel the relay task when the stream terminates by cancellation. The cancellation propagates through the upstream stream's own termination handler to the backend's generation task, releasing the model promptly. Natural completion still finishes the relay and records the transcript entry as before. --- Sources/AnyLanguageModel/LanguageModelSession.swift | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/Sources/AnyLanguageModel/LanguageModelSession.swift b/Sources/AnyLanguageModel/LanguageModelSession.swift index 5ac33376..688a126c 100644 --- a/Sources/AnyLanguageModel/LanguageModelSession.swift +++ b/Sources/AnyLanguageModel/LanguageModelSession.swift @@ -143,7 +143,7 @@ public final class LanguageModelSession: @unchecked Sendable { let session = self let relay = AsyncThrowingStream.Snapshot, any Error> { continuation in let stream = upstream - Task { + let task = Task { session.beginResponding() var lastSnapshot: ResponseStream.Snapshot? do { @@ -178,6 +178,11 @@ public final class LanguageModelSession: @unchecked Sendable { } session.endResponding() } + continuation.onTermination = { termination in + if case .cancelled = termination { + task.cancel() + } + } } return ResponseStream(stream: relay) }