From 76fbc711085924096caea9e332fc9c596f1d296e Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Sun, 26 Apr 2026 15:38:39 -0400
Subject: [PATCH 1/2] [AI] Propagate `Task` cancellation when streaming

---
 FirebaseAI/Sources/Chat.swift                | 5 ++++-
 FirebaseAI/Sources/GenerativeAIService.swift | 5 ++++-
 FirebaseAI/Sources/GenerativeModel.swift     | 5 ++++-
 3 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/FirebaseAI/Sources/Chat.swift b/FirebaseAI/Sources/Chat.swift
index 1f7a7d47b6e..3cca90e6af8 100644
--- a/FirebaseAI/Sources/Chat.swift
+++ b/FirebaseAI/Sources/Chat.swift
@@ -115,7 +115,7 @@ public final class Chat: Sendable {
     let request = history + newContent
     let stream = try model.generateContentStream(request, generationConfig: generationConfig)
     return AsyncThrowingStream { continuation in
-      Task {
+      let task = Task {
         var aggregatedContent: [ModelContent] = []
 
         do {
@@ -142,6 +142,9 @@ public final class Chat: Sendable {
         self._history.append(aggregated)
         continuation.finish()
       }
+      continuation.onTermination = { _ in
+        task.cancel()
+      }
     }
   }
 
diff --git a/FirebaseAI/Sources/GenerativeAIService.swift b/FirebaseAI/Sources/GenerativeAIService.swift
index abeb533404d..2c54671e487 100644
--- a/FirebaseAI/Sources/GenerativeAIService.swift
+++ b/FirebaseAI/Sources/GenerativeAIService.swift
@@ -72,7 +72,7 @@ struct GenerativeAIService {
   func loadRequestStream<T: GenerativeAIRequest>(request: T)
     -> AsyncThrowingStream<T.Response, Error> where T: Sendable {
     return AsyncThrowingStream { continuation in
-      Task {
+      let task = Task {
         let urlRequest: URLRequest
         do {
           urlRequest = try await self.urlRequest(request: request)
@@ -160,6 +160,9 @@ struct GenerativeAIService {
 
         continuation.finish(throwing: nil)
       }
+      continuation.onTermination = { _ in
+        task.cancel()
+      }
     }
   }
 
diff --git a/FirebaseAI/Sources/GenerativeModel.swift b/FirebaseAI/Sources/GenerativeModel.swift
index 20f823b9454..0ef5dc9e821 100644
--- a/FirebaseAI/Sources/GenerativeModel.swift
+++ b/FirebaseAI/Sources/GenerativeModel.swift
@@ -322,7 +322,7 @@ public final class GenerativeModel: Sendable {
 
     return AsyncThrowingStream { continuation in
       let responseStream = generativeAIService.loadRequestStream(request: generateContentRequest)
-      Task {
+      let task = Task {
         do {
           var didYieldResponse = false
           for try await response in responseStream {
@@ -368,6 +368,9 @@ public final class GenerativeModel: Sendable {
           return
         }
       }
+      continuation.onTermination = { _ in
+        task.cancel()
+      }
     }
   }
 

From 3fbc13fd8f6f79bfd88ada77b9657705021ba079 Mon Sep 17 00:00:00 2001
From: Andrew Heard <andrewheard@google.com>
Date: Tue, 30 Jun 2026 15:43:45 -0400
Subject: [PATCH 2/2] Update GenerativeAIService.swift

---
 FirebaseAI/Sources/GenerativeAIService.swift | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/FirebaseAI/Sources/GenerativeAIService.swift b/FirebaseAI/Sources/GenerativeAIService.swift
index 24c539c7aba..73394744258 100644
--- a/FirebaseAI/Sources/GenerativeAIService.swift
+++ b/FirebaseAI/Sources/GenerativeAIService.swift
@@ -137,9 +137,6 @@ struct GenerativeAIService {
       continuation.onTermination = { @Sendable _ in
         task.cancel()
       }
-      continuation.onTermination = { _ in
-        task.cancel()
-      }
     }
   }